HParams for training languagemodel_wikitext103_l4k with memory.
973422
def transformer_wikitext103_l4k_memory_v0():
  """"""
  hparams = transformer_wikitext103_l4k_v0()

  hparams.split_targets_chunk_length = 64
  hparams.split_targets_max_chunks = 64
  hparams.split_targets_strided_training = True
  hparams.add_hparam("memory_type", "transformer_xl")

  # 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))  # 262144

  hparams.pos = None
  hparams.self_attention_type = "dot_product_relative"
  hparams.max_relative_position = 2 * hparams.split_targets_chunk_length

  hparams.add_hparam("unconditional", True)
  hparams.add_hparam("recurrent_memory_batch_size", 0)  # 0 = try to guess
  # By default, cache one chunk only (like Transformer-XL)
  hparams.add_hparam("num_memory_items", hparams.split_targets_chunk_length)

  return hparams
948691
def attention_lm_translation():
  """"""
  hparams = attention_lm_base()
  hparams.layer_preprocess_sequence = "n"
  hparams.layer_postprocess_sequence = "da"
  hparams.learning_rate = 0.4
  hparams.prepend_mode = "prepend_inputs_masked_attention"
  hparams.max_length = 512
  hparams.label_smoothing = 0.1
  hparams.shared_embedding_and_softmax_weights = True
  return hparams
810724
def mtf_bitransformer_base():
  """"""
  hparams = mtf_transformer2_base()
  hparams.max_length = 256
  hparams.shared_embedding = True
  # HYPERPARAMETERS FOR THE LAYER STACKS
  hparams.add_hparam("encoder_layers", ["self_att", "drd"] * 6)
  hparams.add_hparam("decoder_layers", ["self_att", "enc_att", "drd"] * 6)
  hparams.add_hparam("encoder_num_layers", 6)
  hparams.add_hparam("decoder_num_layers", 6)
  # number of heads in multihead attention
  hparams.add_hparam("encoder_num_heads", 8)
  hparams.add_hparam("decoder_num_heads", 8)
  hparams.add_hparam("local_attention_radius", 128)

  # default of 0 for standard transformer behavior
  # 1 means a single set of keys and values that are read by all query heads
  hparams.add_hparam("encoder_num_memory_heads", 0)
  hparams.add_hparam("decoder_num_memory_heads", 0)
  # share attention keys and values
  hparams.add_hparam("encoder_shared_kv", False)
  hparams.add_hparam("decoder_shared_kv", False)

  # Parameters for computing the maximum decode length in beam search.
  # Maximum decode length is:
  #    min(max_length,
  #        decode_length_multiplier * input_length + decode_length_constant)
  hparams.add_hparam("decode_length_multiplier", 1.5)
  hparams.add_hparam("decode_length_constant", 10.0)
  # used during decoding
  hparams.add_hparam("alpha", 0.6)
  hparams.sampling_temp = 0.0
  return hparams
520851
def vqa_recurrent_self_attention_base():
  """"""
  hparams = universal_transformer.universal_transformer_base()
  hparams.batch_size = 1024
  hparams.use_fixed_batch_size = True
  hparams.weight_decay = 0.
  hparams.clip_grad_norm = 0.
  # use default initializer
  # hparams.initializer = "xavier"
  hparams.learning_rate_schedule = (
      "constant*linear_warmup*rsqrt_normalized_decay")
  hparams.learning_rate_warmup_steps = 8000
  hparams.learning_rate_constant = 7e-4
  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.1
  hparams.multiply_embedding_mode = "sqrt_depth"

  # add new hparams
  # use raw image as input
  hparams.add_hparam("image_input_type", "feature")
  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)

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

  # 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")

  # iterative part
  hparams.transformer_ffn_type = "fc"

  return hparams
30448
def data_loader(self, sentences, shuffle=False):
        """"""
        dataset = BertEmbeddingDataset(sentences, self.transform)
        return DataLoader(dataset=dataset, batch_size=self.batch_size, shuffle=shuffle)
762304
def attention_lm_moe_translation():
  """"""
  hparams = attention_lm_moe_base()
  hparams.layer_preprocess_sequence = "n"
  hparams.layer_postprocess_sequence = "da"
  hparams.learning_rate = 0.4
  hparams.prepend_mode = "prepend_inputs_masked_attention"
  hparams.max_length = 512
  hparams.label_smoothing = 0.1
  hparams.layer_prepostprocess_dropout = 0.2
  hparams.num_hidden_layers = 6
  hparams.moe_layers = "0,1,2,3,4,5"
  hparams.shared_embedding_and_softmax_weights = True
  return hparams
1253887
def mtf_image_transformer_base_imagenet_mp():
  """"""
  hparams = mtf_image_transformer_base_imagenet()
  hparams.mesh_shape = "model:4;batch:8"
  hparams.layout = "batch:batch;d_ff:model;heads:model"
  hparams.batch_size = 32
  hparams.num_heads = 8
  hparams.d_ff = 8192
  hparams.learning_rate_warmup_steps = 31250
  hparams.unconditional = True
  return hparams
907460
def mtf_transformer_lm_baseline():
  
  hparams = mtf_transformer_paper_lm(-1)
  hparams.batch_size = 128
  hparams.learning_rate_decay_steps = 27200  # one epoch on lm1b
  hparams.mesh_shape = "batch:8"
  return hparams
570747
def basic_range1(ranged_hparams):
  """"""
  rhp = ranged_hparams
  rhp.set_discrete("batch_size", [1024, 2048, 4096])
  rhp.set_discrete("num_hidden_layers", [1, 2, 3, 4, 5, 6])
  rhp.set_discrete("hidden_size", [32, 64, 128, 256, 512], scale=rhp.LOG_SCALE)
  rhp.set_discrete("kernel_height", [1, 3, 5, 7])
  rhp.set_discrete("kernel_width", [1, 3, 5, 7])
  rhp.set_discrete("compress_steps", [0, 1, 2])
  rhp.set_float("dropout", 0.0, 0.5)
  rhp.set_float("weight_decay", 1e-4, 10.0, scale=rhp.LOG_SCALE)
  rhp.set_float("label_smoothing", 0.0, 0.2)
  rhp.set_float("clip_grad_norm", 0.01, 50.0, scale=rhp.LOG_SCALE)
  rhp.set_float("learning_rate", 0.005, 2.0, scale=rhp.LOG_SCALE)
  rhp.set_categorical("initializer",
                      ["uniform", "orthogonal", "uniform_unit_scaling"])
  rhp.set_float("initializer_gain", 0.5, 3.5)
  rhp.set_categorical("learning_rate_decay_scheme",
                      ["none", "sqrt", "noam", "exp"])
  rhp.set_float("optimizer_adam_epsilon", 1e-7, 1e-2, scale=rhp.LOG_SCALE)
  rhp.set_float("optimizer_adam_beta1", 0.8, 0.9)
  rhp.set_float("optimizer_adam_beta2", 0.995, 0.999)
  rhp.set_categorical(
      "optimizer",
      ["adam", "adagrad", "momentum", "rms_prop", "sgd", "yellow_fin"])
1140599
def self_attention_layer(hparams, prefix):
  """"""
  return transformer_layers.SelfAttention(
      num_heads=hparams.get(prefix + "num_heads"),
      num_memory_heads=hparams.get(prefix + "num_memory_heads"),
      key_value_size=hparams.d_kv,
      shared_kv=hparams.get(prefix + "shared_kv", False),
      attention_kwargs=attention_kwargs_from_hparams(hparams))
245567
def attention_lm_moe_memory_efficient():
  """"""
  hparams = attention_lm_moe_large()
  hparams.diet_experts = True
  hparams.layer_preprocess_sequence = "n"
  hparams.layer_postprocess_sequence = "da"
  hparams.layer_prepostprocess_dropout = 0.0
  hparams.memory_efficient_ffn = True
  hparams.attention_type = AttentionType.MEMORY_EFFICIENT
  hparams.num_heads = 8
  hparams.factored_logits = True
  return hparams
355335
def transformer_cifar10_memory_v0():
  """"""
  hparams = transformer_wikitext103_l4k_memory_v0()

  hparams.num_hidden_layers = 6

  hparams.max_length = 32 * 32 * 3
  hparams.split_targets_chunk_length = 64 * 3
  hparams.split_targets_max_chunks = int(
      hparams.max_length / hparams.split_targets_chunk_length)
  hparams.num_memory_items = 128 * 3

  # Since this is an image problem, batch size refers to examples (not tokens)
  target_images_per_batch = 4
  hparams.batch_size = int(target_images_per_batch * (
      hparams.max_length / hparams.split_targets_chunk_length))

  # The recurrent memory needs to know the actual batch size (in sequences)
  hparams.recurrent_memory_batch_size = hparams.batch_size

  hparams.max_relative_position = (
      hparams.num_memory_items + hparams.split_targets_chunk_length)

  return hparams
1092226
def lmx_moe_h1k_f4k_x32():
  """"""
  hparams = lmx_h1k_f4k()
  hparams.ffn_layer = "local_moe_tpu"
  hparams.moe_num_experts = 32
  hparams.weight_dtype = "bfloat16"
  hparams.batch_size = 8192
  return hparams
996490
def lmx_h4k_f16k():
  """"""
  hparams = lmx_base()
  hparams.hidden_size = 4096
  hparams.filter_size = 16384
  hparams.batch_size = 1024
  hparams.weight_dtype = "bfloat16"
  return hparams
61368
def wiki_2x2_base():
  
  hparams = mtf_transformer.mtf_transformer_base_lm()
  hparams.shared_embedding_and_softmax_weights = False
  # no dropout - dataset is big enough to avoid overfitting.
  hparams.attention_dropout = 0.0
  hparams.relu_dropout = 0.0
  hparams.layer_prepostprocess_dropout = 0.0
  hparams.max_length = 1024
  # 4 sequences per core
  hparams.batch_size = 32
  # We don't use linear decay in these experiments, since we don't want
  # a sharp jump in quality at the end of the training schedule.
  # You can insert this once you find the right architecture.
  hparams.learning_rate_schedule = "rsqrt_decay"
  hparams.mesh_shape = "all:8"
  hparams.layout = "batch:all;experts:all"

  # parameters for mixture-of-experts
  moe.set_default_moe_hparams(hparams)
  hparams.moe_num_experts = 16
  hparams.moe_hidden_size = 8192

  hparams.decoder_layers = ["att", "drd"] * 6
  hparams.d_model = 1024
  hparams.d_ff = 2048
  hparams.d_kv = 128
  hparams.num_heads = 4

  return hparams
1053406
def super_lm_base():
  """"""
  hparams = common_hparams.basic_params1()
  hparams.hidden_size = 512
  hparams.moe_hidden_sizes = "512"
  hparams.batch_size = 16384
  hparams.max_length = 0
  # All hyperparameters ending in "dropout" are automatically set to 0.0
  # when not in training mode.
  hparams.layer_prepostprocess_dropout = 0.0
  hparams.symbol_dropout = 0.1
  hparams.add_hparam("attention_dropout", 0.0)
  hparams.label_smoothing = 0.0
  hparams.clip_grad_norm = 0.  # i.e. no gradient clipping
  hparams.optimizer = "Adafactor"
  hparams.learning_rate_decay_scheme = "noam"
  hparams.learning_rate = 0.1
  hparams.learning_rate_warmup_steps = 8000
  hparams.initializer_gain = 1.0
  hparams.initializer = "uniform_unit_scaling"
  hparams.weight_decay = 0.0
  hparams.shared_embedding_and_softmax_weights = False
  hparams.layer_preprocess_sequence = "n"
  hparams.layer_postprocess_sequence = "da"
  # we only want one data shard.
  hparams.no_data_parallelism = True
  # bypass the symbol modality so that we can use model parallelism.
  hparams.bottom = {
      "inputs": modalities.identity_bottom,
      "targets": modalities.identity_bottom,
  }
  hparams.top = {
      "targets": modalities.identity_top,
  }
  hparams.add_hparam("filter_size", 512)
  hparams.add_hparam("mix_fraction", 0.5)
  # attention-related flags
  hparams.add_hparam("multihead_attention_num_heads", 4)
  hparams.add_hparam("multihead_attention_key_channels", 0)
  hparams.add_hparam("multihead_attention_value_channels", 0)
  hparams.add_hparam("pos", "timing")  # timing, none
  hparams.add_hparam(
      "layers", ("n,att,m,d,a," "n,ffn,m,d,a,") * 4 + "n,ffn,d")
  # Number of model shards - each one has separate parameters.
  # Changing this number invalidates checkpoints.
  hparams.add_hparam("num_model_shards", 8)
  hparams.add_hparam("diet_experts", False)
  return hparams
9241
def aggregate_task_lm_losses(hparams,
                             problem_hparams,
                             logits,
                             feature_name,
                             feature):
  """"""
  summaries = []
  vocab_size = problem_hparams.vocab_size[feature_name]
  if vocab_size is not None and hasattr(hparams, "vocab_divisor"):
    vocab_size += (-vocab_size) % hparams.vocab_divisor
  modality = problem_hparams.modality[feature_name]
  loss = hparams.loss.get(feature_name, modalities.get_loss(modality))
  weights_fn = hparams.weights_fn.get(
      feature_name, modalities.get_weights_fn(modality))
  loss_num = 0.
  loss_den = 0.
  for task in hparams.problem.task_list:
    loss_num_, loss_den_ = loss(
        logits, feature,
        lambda x: common_layers.weights_multi_problem_all(x, task.task_id),  # pylint: disable=cell-var-from-loop
        hparams, vocab_size, weights_fn)

    loss_num += loss_num_
    loss_den += loss_den_

    loss_val = loss_num_ / tf.maximum(1.0, loss_den_)
    summaries.append([task.name+"_loss", loss_val])

  return loss_num, loss_den, summaries