Hparams for machine translation with ~1.1B parameters.
5775
def transformer_tpu_1b():
  """"""
  hparams = transformer_tpu()
  hparams.hidden_size = 2048
  hparams.filter_size = 8192
  hparams.num_hidden_layers = 8
  # smaller batch size to avoid OOM
  hparams.batch_size = 1024
  hparams.activation_dtype = "bfloat16"
  hparams.weight_dtype = "bfloat16"
  # maximize number of parameters relative to computation by not sharing.
  hparams.shared_embedding_and_softmax_weights = False
  return hparams
213331
def transform(instance, tokenizer, max_seq_length, max_predictions_per_seq, do_pad=True):
    """"""
    pad = tokenizer.convert_tokens_to_ids(['[PAD]'])[0]
    input_ids = tokenizer.convert_tokens_to_ids(instance.tokens)
    input_mask = [1] * len(input_ids)
    segment_ids = list(instance.segment_ids)
    assert len(input_ids) <= max_seq_length
    valid_lengths = len(input_ids)

    masked_lm_positions = list(instance.masked_lm_positions)
    masked_lm_ids = tokenizer.convert_tokens_to_ids(
        instance.masked_lm_labels)
    masked_lm_weights = [1.0] * len(masked_lm_ids)
    masked_lm_valid_lengths = len(masked_lm_ids)
    if do_pad:
        while len(input_ids) < max_seq_length:
            input_ids.append(pad)
            # Padding index MUST be defined to 0 on input_mask, segment_ids
            input_mask.append(0)
            segment_ids.append(0)
        while len(masked_lm_positions) < max_predictions_per_seq:
            masked_lm_positions.append(0)
            masked_lm_ids.append(pad)
            masked_lm_weights.append(0.0)

        assert len(input_ids) == max_seq_length
        assert len(input_mask) == max_seq_length
        assert len(segment_ids) == max_seq_length

    next_sentence_label = 1 if instance.is_random_next else 0

    features = {}
    features['input_ids'] = input_ids
    features['input_mask'] = input_mask
    features['segment_ids'] = segment_ids
    features['masked_lm_positions'] = masked_lm_positions
    features['masked_lm_ids'] = masked_lm_ids
    features['segment_a_lengths'] = [instance.segment_a_lengths]
    features['segment_b_lengths'] = [instance.segment_b_lengths]
    features['masked_lm_ids'] = masked_lm_ids
    features['masked_lm_weights'] = masked_lm_weights
    features['next_sentence_labels'] = [next_sentence_label]
    features['valid_lengths'] = [valid_lengths]
    features['masked_lm_valid_lengths'] = [masked_lm_valid_lengths]

    return features
214995
def autoencoder_ordered_text_small():
  """"""
  hparams = autoencoder_ordered_text()
  hparams.bottleneck_bits = 32
  hparams.num_hidden_layers = 3
  hparams.hidden_size = 64
  hparams.max_hidden_size = 512
  hparams.bottleneck_noise = 0.0
  hparams.autoregressive_mode = "conv5"
  hparams.sample_height = 4
  return hparams
1309195
def mtf_transformer_tiny():
  """"""
  hparams = mtf_transformer_base()
  hparams.d_model = 128
  hparams.d_ff = 512
  hparams.batch_size = 8
  hparams.encoder_layers = ["att", "drd"] * 2
  hparams.decoder_layers = ["att", "enc_att", "drd"] * 2
  hparams.num_heads = 8
  # data parallelism and model-parallelism
  hparams.mesh_shape = "batch:2;model:4"
  hparams.activation_dtype = "float32"
  return hparams
653384
def autoencoder_basic():
  """"""
  hparams = common_hparams.basic_params1()
  hparams.optimizer = "adam"
  hparams.learning_rate_constant = 0.0002
  hparams.learning_rate_warmup_steps = 500
  hparams.learning_rate_schedule = "constant * linear_warmup"
  hparams.label_smoothing = 0.0
  hparams.batch_size = 128
  hparams.hidden_size = 64
  hparams.num_hidden_layers = 5
  hparams.initializer = "uniform_unit_scaling"
  hparams.initializer_gain = 1.0
  hparams.weight_decay = 0.0
  hparams.kernel_height = 4
  hparams.kernel_width = 4
  hparams.dropout = 0.05
  hparams.add_hparam("max_hidden_size", 1024)
  hparams.add_hparam("bottleneck_bits", 128)
  hparams.add_hparam("bottleneck_shared_bits", 0)
  hparams.add_hparam("bottleneck_shared_bits_start_warmup", 0)
  hparams.add_hparam("bottleneck_shared_bits_stop_warmup", 0)
  hparams.add_hparam("bottleneck_noise", 0.1)
  hparams.add_hparam("bottleneck_warmup_steps", 2000)
  hparams.add_hparam("sample_height", 32)
  hparams.add_hparam("sample_width", 32)
  hparams.add_hparam("discriminator_batchnorm", True)
  hparams.add_hparam("num_sliced_vecs", 20000)
  hparams.add_hparam("sliced_do_tanh", int(True))
  hparams.add_hparam("discriminator_size", 256)
  hparams.add_hparam("discriminator_kernel_size", 6)
  hparams.add_hparam("discriminator_strides", 4)
  hparams.add_hparam("discriminator_pure_mean", int(False))
  hparams.add_hparam("code_loss_factor", 1.0)
  hparams.add_hparam("gan_codes_warmup_steps", 16000)
  hparams.add_hparam("gan_loss_factor", 0.0)
  hparams.add_hparam("bottleneck_l2_factor", 0.05)
  hparams.add_hparam("gumbel_temperature", 0.5)
  hparams.add_hparam("gumbel_noise_factor", 0.5)
  hparams.add_hparam("vq_temperature", 0.001)
  hparams.add_hparam("use_vq_loss", int(False))
  hparams.add_hparam("discriminator", "double")
  return hparams
64935
def transformer_moe_prepend_8k():
  """"""
  hparams = transformer_moe_8k()
  hparams.prepend_mode = "prepend_inputs_masked_attention"
  hparams.eval_drop_long_sequences = False
  hparams.max_input_seq_length = 7500
  hparams.default_ff = "sepm"
  hparams.layer_types = "locm/redm/locm-moe/redm/locm"
  hparams.moe_num_experts = 256
  return hparams
82338
def vqa_self_attention_base():
  """"""
  hparams = common_hparams.basic_params1()
  hparams.batch_size = 128
  hparams.use_fixed_batch_size = True,
  hparams.optimizer = "adam"
  hparams.optimizer_adam_beta1 = 0.9
  hparams.optimizer_adam_beta2 = 0.997
  hparams.optimizer_adam_epsilon = 1e-9
  hparams.weight_decay = 0.
  hparams.clip_grad_norm = 0.
  hparams.initializer = "xavier"
  hparams.learning_rate_schedule = (
      "constant*linear_warmup*rsqrt_normalized_decay")
  hparams.learning_rate_warmup_steps = 8000
  hparams.learning_rate_constant = 1e-3
  hparams.learning_rate_decay_rate = 0.5
  hparams.learning_rate_decay_steps = 50000
  hparams.dropout = 0.5
  hparams.summarize_grads = True
  hparams.summarize_vars = True

  # not used hparams
  hparams.label_smoothing = 0.
  hparams.multiply_embedding_mode = "sqrt_depth"

  # add new hparams
  # use raw image as input
  hparams.add_hparam("image_input_type", "image")
  hparams.add_hparam("image_model_fn", "resnet_v1_152")
  hparams.add_hparam("resize_side", 512)
  hparams.add_hparam("height", 448)
  hparams.add_hparam("width", 448)
  hparams.add_hparam("distort", True)
  hparams.add_hparam("train_resnet", False)

  # image parts
  hparams.add_hparam("image_feat_preprocess_proj", True)
  hparams.add_hparam("image_feat_preprocess_layernorm", True)
  hparams.add_hparam("image_feat_encode", True)
  hparams.add_hparam("image_hidden_size", 0)  # default to hidden_size
  hparams.add_hparam("image_filter_size", 0)  # defaults to filter_size

  # question hidden size
  hparams.hidden_size = 512
  hparams.filter_size = 1024
  hparams.num_hidden_layers = 4

  hparams.add_hparam("multimodal_combine", "concat")
  hparams.add_hparam("num_mlp_layers", 1)
  hparams.add_hparam("mlp_size", 1024)

  # self attention parts
  hparams.norm_type = "layer"
  hparams.layer_preprocess_sequence = "n"
  hparams.layer_postprocess_sequence = "da"
  hparams.layer_prepostprocess_dropout = 0.1
  hparams.attention_dropout = 0.1
  hparams.relu_dropout = 0.1
  hparams.add_hparam("pos", "timing")
  hparams.add_hparam("num_encoder_layers", 0)
  hparams.add_hparam("num_decoder_layers", 0)
  hparams.add_hparam("num_heads", 8)
  hparams.add_hparam("attention_key_channels", 0)
  hparams.add_hparam("attention_value_channels", 0)
  hparams.add_hparam("self_attention_type", "dot_product")
  hparams.add_hparam("image_self_attention_type", "dot_product")
  hparams.add_hparam("question_self_attention_type", "dot_product")
  hparams.add_hparam("block_length", 1)
  hparams.add_hparam("scale_dotproduct", True)

  # iterative part
  hparams.add_hparam("num_rec_steps", 3)

  return hparams
478750
def mtr_lm_v1():
  
  hparams = mtr_lm_dense(0)
  hparams.layers = (["local_self_att", "local_self_att", "drd",
                     "self_att", "drd", "local_self_att",
                     "local_self_att", "moe_2d"] * 4)[:-1]
  hparams.d_kv = 128
  hparams.moe_expert_x = 8
  hparams.moe_expert_y = 4
  hparams.moe_hidden_size = 32768
  hparams.d_ff = 2048
  hparams.num_memory_heads = 0
  hparams.mesh_shape = "b0:4;b1:8"
  hparams.layout = "outer_batch:b0;inner_batch:b1,expert_x:b1,expert_y:b0"
  hparams.outer_batch_size = 4
  return hparams
299140
def transformer_wikitext103_l16k_memory_v0():
  """"""
  hparams = transformer_wikitext103_l4k_memory_v0()

  hparams.max_length = 16384
  hparams.split_targets_chunk_length = 64
  hparams.split_targets_max_chunks = int(
      hparams.max_length / hparams.split_targets_chunk_length)

  # The hparams specify batch size *before* chunking, but we want to have a
  # consistent 4K batch size *after* chunking to fully utilize the hardware.
  target_tokens_per_batch = 4096
  hparams.batch_size = int(target_tokens_per_batch * (
      hparams.max_length / hparams.split_targets_chunk_length))

  hparams.max_relative_position = 2 * hparams.split_targets_chunk_length

  return hparams
1064209
def img2img_transformer2d_n31():
  """"""
  hparams = img2img_transformer2d_base()
  hparams.batch_size = 1
  hparams.num_encoder_layers = 6
  hparams.num_decoder_layers = 12
  hparams.num_heads = 8
  hparams.query_shape = (16, 32)
  hparams.memory_flange = (16, 32)
  return hparams
875173
def mtr_tr_dense_local(sz):
  """"""
  hparams = mtr_tr_dense(sz)
  hparams.decoder_layers = ["local_self_att", "enc_att", "drd"] * 6
  hparams.local_attention_radius = 32
  return hparams
1131142
def transformer_base_v1():
  """"""
  hparams = common_hparams.basic_params1()
  hparams.norm_type = "layer"
  hparams.hidden_size = 512
  hparams.batch_size = 4096
  hparams.max_length = 256
  hparams.clip_grad_norm = 0.  # i.e. no gradient clipping
  hparams.optimizer_adam_epsilon = 1e-9
  hparams.learning_rate_schedule = "legacy"
  hparams.learning_rate_decay_scheme = "noam"
  hparams.learning_rate = 0.1
  hparams.learning_rate_warmup_steps = 4000
  hparams.initializer_gain = 1.0
  hparams.num_hidden_layers = 6
  hparams.initializer = "uniform_unit_scaling"
  hparams.weight_decay = 0.0
  hparams.optimizer_adam_beta1 = 0.9
  hparams.optimizer_adam_beta2 = 0.98
  hparams.num_sampled_classes = 0
  hparams.label_smoothing = 0.1
  hparams.shared_embedding_and_softmax_weights = True
  hparams.symbol_modality_num_shards = 16

  # Add new ones like this.
  hparams.add_hparam("filter_size", 2048)
  # Layer-related flags. If zero, these fall back on hparams.num_hidden_layers.
  hparams.add_hparam("num_encoder_layers", 0)
  hparams.add_hparam("num_decoder_layers", 0)
  # Attention-related flags.
  hparams.add_hparam("num_heads", 8)
  hparams.add_hparam("attention_key_channels", 0)
  hparams.add_hparam("attention_value_channels", 0)
  hparams.add_hparam("ffn_layer", "dense_relu_dense")
  hparams.add_hparam("parameter_attention_key_channels", 0)
  hparams.add_hparam("parameter_attention_value_channels", 0)
  # All hyperparameters ending in "dropout" are automatically set to 0.0
  # when not in training mode.
  hparams.add_hparam("attention_dropout", 0.0)
  hparams.add_hparam("attention_dropout_broadcast_dims", "")
  hparams.add_hparam("relu_dropout", 0.0)
  hparams.add_hparam("relu_dropout_broadcast_dims", "")
  hparams.add_hparam("pos", "timing")  # timing, none
  hparams.add_hparam("nbr_decoder_problems", 1)
  hparams.add_hparam("proximity_bias", False)
  hparams.add_hparam("causal_decoder_self_attention", True)
  hparams.add_hparam("use_pad_remover", True)
  hparams.add_hparam("self_attention_type", "dot_product")
  hparams.add_hparam("conv_first_kernel", 3)
  hparams.add_hparam("attention_variables_3d", False)
  hparams.add_hparam("use_target_space_embedding", True)
  # These parameters are only used when ffn_layer=="local_moe_tpu"
  hparams.add_hparam("moe_overhead_train", 1.0)
  hparams.add_hparam("moe_overhead_eval", 2.0)
  hparams.moe_num_experts = 16
  hparams.moe_loss_coef = 1e-3
  # If specified, use this value instead of problem name in metrics.py.
  # This is useful for programs that can automatically compare experiments side
  #   by side based on the same metric names.
  hparams.add_hparam("overload_eval_metric_name", "")
  # For making a transformer encoder unidirectional by using masked
  # attention.
  hparams.add_hparam("unidirectional_encoder", False)
  # For hard attention.
  hparams.add_hparam("hard_attention_k", 0)
  return hparams
639701
def preprocess_data(tokenizer, task, batch_size, dev_batch_size, max_len, pad=False):
    """"""
    # transformation
    trans = BERTDatasetTransform(
        tokenizer,
        max_len,
        labels=task.get_labels(),
        pad=pad,
        pair=task.is_pair,
        label_dtype='float32' if not task.get_labels() else 'int32')

    data_train = task('train').transform(trans, lazy=False)
    data_train_len = data_train.transform(
        lambda input_id, length, segment_id, label_id: length)

    num_samples_train = len(data_train)
    # bucket sampler
    batchify_fn = nlp.data.batchify.Tuple(
        nlp.data.batchify.Pad(axis=0), nlp.data.batchify.Stack(),
        nlp.data.batchify.Pad(axis=0),
        nlp.data.batchify.Stack(
            'float32' if not task.get_labels() else 'int32'))
    batch_sampler = nlp.data.sampler.FixedBucketSampler(
        data_train_len,
        batch_size=batch_size,
        num_buckets=10,
        ratio=0,
        shuffle=True)
    # data loaders
    dataloader_train = gluon.data.DataLoader(
        dataset=data_train,
        num_workers=1,
        batch_sampler=batch_sampler,
        batchify_fn=batchify_fn)
    if task.task_name == 'MNLI':
        data_dev_matched = task('dev_matched').transform(trans, lazy=False)
        data_dev_mismatched = task('dev_mismatched').transform(trans, lazy=False)

        dataloader_dev_matched = mx.gluon.data.DataLoader(
            data_dev_matched, batch_size=dev_batch_size,
            num_workers=1, shuffle=False, batchify_fn=batchify_fn)
        dataloader_dev_mismatched = mx.gluon.data.DataLoader(
            data_dev_mismatched, batch_size=dev_batch_size,
            num_workers=1, shuffle=False, batchify_fn=batchify_fn)
        return dataloader_train, dataloader_dev_matched, \
            dataloader_dev_mismatched, num_samples_train
    else:
        data_dev = task('dev').transform(trans, lazy=False)
        dataloader_dev = mx.gluon.data.DataLoader(
            data_dev,
            batch_size=dev_batch_size,
            num_workers=1,
            shuffle=False,
            batchify_fn=batchify_fn)
        return dataloader_train, dataloader_dev, num_samples_train
616659
def transformer_lm_tpu_0():
  """"""
  hparams = transformer_clean_big()
  update_hparams_for_tpu(hparams)
  hparams.num_heads = 4  # Heads are expensive on TPUs.
  hparams.batch_size = 4096
  hparams.shared_embedding_and_softmax_weights = False
  hparams.layer_prepostprocess_dropout = 0.1
  return hparams
599364
def transformer_librispeech_v2():
  """"""
  hparams = transformer_base()

  hparams.max_length = 1240000
  hparams.max_input_seq_length = 1550
  hparams.max_target_seq_length = 350
  hparams.batch_size = 16
  hparams.num_decoder_layers = 4
  hparams.num_encoder_layers = 6
  hparams.hidden_size = 384
  hparams.learning_rate = 0.15
  hparams.daisy_chain_variables = False
  hparams.filter_size = 1536
  hparams.num_heads = 2
  hparams.ffn_layer = "conv_relu_conv"
  hparams.conv_first_kernel = 9
  hparams.weight_decay = 0
  hparams.layer_prepostprocess_dropout = 0.2
  hparams.relu_dropout = 0.2

  return hparams
851590
def mimic_adam_with_adafactor(hparams):
  
  assert "adam" in hparams.optimizer
  hparams.optimizer = "adafactor"
  hparams.optimizer_adafactor_beta1 = hparams.optimizer_adam_beta1
  hparams.optimizer_adafactor_beta2 = hparams.optimizer_adam_beta2
  hparams.optimizer_adafactor_multiply_by_parameter_scale = False
  hparams.optimizer_adafactor_factored = False
  hparams.optimizer_adafactor_clipping_threshold = None
  hparams.optimizer_adafactor_decay_type = "adam"
1200091
def _copy_problem_hparams(p_hparams):
  """"""
  p = p_hparams
  # Duplicate input modality.
  p.modality["targets"] = p.modality["inputs"]
  # Duplicate input vocab size.
  p.vocab_size["targets"] = p.vocab_size["inputs"]
  # Duplicate input vocabulary.
  p.vocabulary["targets"] = p.vocabulary["inputs"]
  # Duplicate input space ids.
  p.target_space_id = p.input_space_id
  # Mark that p was reversed.
  p.was_copy = True