Residual discrete autoencoder model, big version.
187739
def autoencoder_residual_discrete_big():
  """"""
  hparams = autoencoder_residual_discrete()
  hparams.hidden_size = 128
  hparams.max_hidden_size = 4096
  hparams.bottleneck_noise = 0.1
  hparams.residual_dropout = 0.4
  return hparams
1211934
def dueling_model(img_in, num_actions, scope, noisy=False, reuse=False,
                  concat_softmax=False):
  """"""
  with tf.variable_scope(scope, reuse=reuse):
    out = img_in
    with tf.variable_scope("convnet"):
      # original architecture
      out = layers.convolution2d(out, num_outputs=32, kernel_size=8,
                                 stride=4, activation_fn=tf.nn.relu)
      out = layers.convolution2d(out, num_outputs=64, kernel_size=4,
                                 stride=2, activation_fn=tf.nn.relu)
      out = layers.convolution2d(out, num_outputs=64, kernel_size=3,
                                 stride=1, activation_fn=tf.nn.relu)
    out = layers.flatten(out)

    with tf.variable_scope("state_value"):
      if noisy:
        # Apply noisy network on fully connected layers
        # ref: https://arxiv.org/abs/1706.10295
        state_hidden = noisy_dense(out, name='noisy_fc1', size=512,
                                   activation_fn=tf.nn.relu)
        state_score = noisy_dense(state_hidden, name='noisy_fc2',
                                  size=1)
      else:
        state_hidden = layers.fully_connected(
            out,
            num_outputs=512,
            activation_fn=tf.nn.relu
        )
        state_score = layers.fully_connected(state_hidden,
                                             num_outputs=1,
                                             activation_fn=None)
    with tf.variable_scope("action_value"):
      if noisy:
        # Apply noisy network on fully connected layers
        # ref: https://arxiv.org/abs/1706.10295
        actions_hidden = noisy_dense(out, name='noisy_fc1', size=512,
                                     activation_fn=tf.nn.relu)
        action_scores = noisy_dense(actions_hidden, name='noisy_fc2',
                                    size=num_actions)
      else:
        actions_hidden = layers.fully_connected(
            out,
            num_outputs=512,
            activation_fn=tf.nn.relu
        )
        action_scores = layers.fully_connected(
            actions_hidden,
            num_outputs=num_actions,
            activation_fn=None
        )
      action_scores_mean = tf.reduce_mean(action_scores, 1)
      action_scores = action_scores - tf.expand_dims(
          action_scores_mean,
          1
      )

    return state_score + action_scores
210279
def conv2d(x_input, w_matrix):
    """"""
    return tf.nn.conv2d(x_input, w_matrix, strides=[1, 1, 1, 1], padding='SAME')
89108
def vq_discrete_unbottleneck(x, hidden_size):
  """"""
  x_shape = common_layers.shape_list(x)
  x = tf.to_float(x)
  bottleneck_size = common_layers.shape_list(x)[-1]
  means, _, _ = get_vq_codebook(bottleneck_size, hidden_size)
  result = tf.matmul(tf.reshape(x, [-1, x_shape[-1]]), means)
  return tf.reshape(result, x_shape[:-1] + [hidden_size])
722209
def general_conv(x,
                 num_filters=64,
                 filter_size=7,
                 stride=1,
                 stddev=0.02,
                 padding="VALID",
                 name="conv",
                 do_norm="instance",
                 do_relu=True,
                 relufactor=0):
  """"""
  with tf.variable_scope(name):
    x = layers().Conv2D(
        num_filters,
        filter_size,
        stride,
        padding,
        activation=None,
        kernel_initializer=tf.truncated_normal_initializer(stddev=stddev),
        bias_initializer=tf.constant_initializer(0.0))(x)
    if do_norm == "layer":
      x = layer_norm(x)
    elif do_norm == "instance":
      x = instance_norm(x)

    if do_relu:
      if relufactor == 0:
        x = tf.nn.relu(x, "relu")
      else:
        x = lrelu(x, leak=relufactor)

    return x
1102506
def residual_conv(x, repeat, k, hparams, name, reuse=None):
  """"""
  with tf.variable_scope(name, reuse=reuse):
    dilations_and_kernels = [((1, 1), k) for _ in range(3)]
    for i in range(repeat):
      with tf.variable_scope("repeat_%d" % i):
        y = common_layers.conv_block(
            common_layers.layer_norm(x, hparams.hidden_size, name="lnorm"),
            hparams.hidden_size,
            dilations_and_kernels,
            padding="SAME",
            name="residual_conv")
        y = tf.nn.dropout(y, 1.0 - hparams.dropout)
        x += y
    return x
191540
def conv(self,
             num_out_channels,
             k_height,
             k_width,
             d_height=1,
             d_width=1,
             mode="SAME",
             input_layer=None,
             num_channels_in=None,
             use_batch_norm=None,
             stddev=None,
             activation="relu",
             bias=0.0):
        """"""
        if input_layer is None:
            input_layer = self.top_layer
        if num_channels_in is None:
            num_channels_in = self.top_size
        kernel_initializer = None
        if stddev is not None:
            kernel_initializer = tf.truncated_normal_initializer(stddev=stddev)
        name = "conv" + str(self.counts["conv"])
        self.counts["conv"] += 1
        with tf.variable_scope(name):
            strides = [1, d_height, d_width, 1]
            if self.data_format == "NCHW":
                strides = [strides[0], strides[3], strides[1], strides[2]]
            if mode != "SAME_RESNET":
                conv = self._conv2d_impl(
                    input_layer,
                    num_channels_in,
                    num_out_channels,
                    kernel_size=[k_height, k_width],
                    strides=[d_height, d_width],
                    padding=mode,
                    kernel_initializer=kernel_initializer)
            else:  # Special padding mode for ResNet models
                if d_height == 1 and d_width == 1:
                    conv = self._conv2d_impl(
                        input_layer,
                        num_channels_in,
                        num_out_channels,
                        kernel_size=[k_height, k_width],
                        strides=[d_height, d_width],
                        padding="SAME",
                        kernel_initializer=kernel_initializer)
                else:
                    rate = 1  # Unused (for 'a trous' convolutions)
                    kernel_height_effective = k_height + (k_height - 1) * (
                        rate - 1)
                    pad_h_beg = (kernel_height_effective - 1) // 2
                    pad_h_end = kernel_height_effective - 1 - pad_h_beg
                    kernel_width_effective = k_width + (k_width - 1) * (
                        rate - 1)
                    pad_w_beg = (kernel_width_effective - 1) // 2
                    pad_w_end = kernel_width_effective - 1 - pad_w_beg
                    padding = [[0, 0], [pad_h_beg, pad_h_end],
                               [pad_w_beg, pad_w_end], [0, 0]]
                    if self.data_format == "NCHW":
                        padding = [
                            padding[0], padding[3], padding[1], padding[2]
                        ]
                    input_layer = tf.pad(input_layer, padding)
                    conv = self._conv2d_impl(
                        input_layer,
                        num_channels_in,
                        num_out_channels,
                        kernel_size=[k_height, k_width],
                        strides=[d_height, d_width],
                        padding="VALID",
                        kernel_initializer=kernel_initializer)
            if use_batch_norm is None:
                use_batch_norm = self.use_batch_norm
            if not use_batch_norm:
                if bias is not None:
                    biases = self.get_variable(
                        "biases", [num_out_channels],
                        self.variable_dtype,
                        self.dtype,
                        initializer=tf.constant_initializer(bias))
                    biased = tf.reshape(
                        tf.nn.bias_add(
                            conv, biases, data_format=self.data_format),
                        conv.get_shape())
                else:
                    biased = conv
            else:
                self.top_layer = conv
                self.top_size = num_out_channels
                biased = self.batch_norm(**self.batch_norm_config)
            if activation == "relu":
                conv1 = tf.nn.relu(biased)
            elif activation == "linear" or activation is None:
                conv1 = biased
            elif activation == "tanh":
                conv1 = tf.nn.tanh(biased)
            else:
                raise KeyError("Invalid activation type \"%s\"" % activation)
            self.top_layer = conv1
            self.top_size = num_out_channels
            return conv1
1036976
def transformer_decoder_layers(inputs,
                               encoder_output,
                               num_layers,
                               hparams,
                               self_attention_bias=None,
                               encoder_decoder_attention_bias=None,
                               attention_type=AttentionType.LOCAL_2D,
                               losses=None,
                               name="transformer"):
  """"""
  x = inputs
  x = tf.nn.dropout(x, 1.0 - hparams.layer_prepostprocess_dropout)
  if attention_type == AttentionType.DILATED:
    assert len(hparams.gap_sizes) == num_layers
  for layer in range(num_layers):
    with tf.variable_scope("%s_layer_%d" % (name, layer)):
      # self-attention + skip connections
      if attention_type == AttentionType.LOCAL_2D:
        y = local_attention_2d(common_layers.layer_preprocess(x, hparams),
                               hparams,
                               attention_type="masked_local_attention_2d")
      elif attention_type == AttentionType.LOCAL_1D:
        y = local_attention_1d(common_layers.layer_preprocess(x, hparams),
                               hparams,
                               attention_type="local_mask_right",
                               q_padding="LEFT", kv_padding="LEFT")
      elif attention_type == AttentionType.RELATIVE_LOCAL_1D:
        y = local_attention_1d(
            common_layers.layer_preprocess(x, hparams),
            hparams,
            attention_type="local_relative_mask_right",
            q_padding="LEFT",
            kv_padding="LEFT")
      elif attention_type == AttentionType.NON_CAUSAL_1D:
        y = local_attention_1d(common_layers.layer_preprocess(x, hparams),
                               hparams,
                               attention_type="local_unmasked",
                               q_padding="VALID", kv_padding="VALID")
      elif attention_type == AttentionType.LOCAL_BLOCK:
        y = local_within_block_attention(
            common_layers.layer_preprocess(x, hparams),
            self_attention_bias, hparams,
            attention_type="local_within_block_mask_right",
            q_padding="LEFT", kv_padding="LEFT")
      elif attention_type == AttentionType.GLOCAL:
        y = local_global_attention(common_layers.layer_preprocess(x, hparams),
                                   self_attention_bias, hparams,
                                   q_padding="LEFT", kv_padding="LEFT")
      elif attention_type == AttentionType.DILATED:
        y = dilated_attention_1d(common_layers.layer_preprocess(x, hparams),
                                 hparams, q_padding="LEFT",
                                 kv_padding="LEFT",
                                 gap_size=hparams.gap_sizes[layer])
      elif attention_type == AttentionType.GLOBAL:
        y = full_self_attention(common_layers.layer_preprocess(x, hparams),
                                self_attention_bias, hparams,
                                q_padding="LEFT", kv_padding="LEFT")
      x = common_layers.layer_postprocess(x, y, hparams)
      # enc-dec attention + skip connections
      if encoder_output is not None:
        y = encdec_attention_1d(common_layers.layer_preprocess(x, hparams),
                                encoder_output,
                                encoder_decoder_attention_bias,
                                hparams)
        x = common_layers.layer_postprocess(x, y, hparams)
      # feed-fwd layers + skip connections
      y = ffn_layer(common_layers.layer_preprocess(x, hparams), hparams,
                    losses=losses)
      x = common_layers.layer_postprocess(x, y, hparams)
  return common_layers.layer_preprocess(x, hparams)
922793
def image_embedding(images,
                    model_fn=resnet_v1_152,
                    trainable=True,
                    is_training=True,
                    weight_decay=0.0001,
                    batch_norm_decay=0.997,
                    batch_norm_epsilon=1e-5,
                    batch_norm_scale=True,
                    add_summaries=False,
                    reuse=False):
  """"""

  is_resnet_training = trainable and is_training

  batch_norm_params = {
      "is_training": is_resnet_training,
      "trainable": trainable,
      "decay": batch_norm_decay,
      "epsilon": batch_norm_epsilon,
      "scale": batch_norm_scale,
  }

  if trainable:
    weights_regularizer = tf.contrib.layers.l2_regularizer(weight_decay)
  else:
    weights_regularizer = None

  with tf.variable_scope(model_fn.__name__, [images], reuse=reuse) as scope:
    with slim.arg_scope(
        [slim.conv2d],
        weights_regularizer=weights_regularizer,
        trainable=trainable):
      with slim.arg_scope(
          [slim.conv2d],
          weights_initializer=slim.variance_scaling_initializer(),
          activation_fn=tf.nn.relu,
          normalizer_fn=slim.batch_norm,
          normalizer_params=batch_norm_params):
        with slim.arg_scope([slim.batch_norm],
                            is_training=is_resnet_training,
                            trainable=trainable):
          with slim.arg_scope([slim.max_pool2d], padding="SAME"):
            net, end_points = model_fn(
                images, num_classes=None, global_pool=False,
                is_training=is_resnet_training,
                reuse=reuse, scope=scope)

  if add_summaries:
    for v in end_points.values():
      tf.contrib.layers.summaries.summarize_activation(v)

  return net
456093
def conv_bn_lrelu(ni:int, nf:int, ks:int=3, stride:int=1)->nn.Sequential:
    ""
    return nn.Sequential(
        nn.Conv2d(ni, nf, kernel_size=ks, bias=False, stride=stride, padding=ks//2),
        nn.BatchNorm2d(nf),
        nn.LeakyReLU(negative_slope=0.1, inplace=True))
1296160
def conv_hidden_relu(inputs,
                     hidden_size,
                     output_size,
                     kernel_size=(1, 1),
                     second_kernel_size=(1, 1),
                     dropout=0.0,
                     **kwargs):
  """"""
  name = kwargs.pop("name") if "name" in kwargs else None
  with tf.variable_scope(name, "conv_hidden_relu", [inputs]):
    if inputs.get_shape().ndims == 3:
      is_3d = True
      inputs = tf.expand_dims(inputs, 2)
    else:
      is_3d = False
    conv_f1 = conv if kernel_size == (1, 1) else separable_conv
    h = conv_f1(
        inputs,
        hidden_size,
        kernel_size,
        activation=tf.nn.relu,
        name="conv1",
        **kwargs)
    if dropout != 0.0:
      h = tf.nn.dropout(h, 1.0 - dropout)
    conv_f2 = conv if second_kernel_size == (1, 1) else separable_conv
    ret = conv_f2(h, output_size, second_kernel_size, name="conv2", **kwargs)
    if is_3d:
      ret = tf.squeeze(ret, 2)
    return ret
1123408
def autoencoder_range(rhp):
  """"""
  rhp.set_float("dropout", 0.01, 0.3)
  rhp.set_float("gan_loss_factor", 0.01, 0.1)
  rhp.set_float("bottleneck_l2_factor", 0.001, 0.1, scale=rhp.LOG_SCALE)
  rhp.set_discrete("bottleneck_warmup_steps", [200, 2000])
  rhp.set_float("gumbel_temperature", 0, 1)
  rhp.set_float("gumbel_noise_factor", 0, 0.5)
845367
def imagetransformerpp_base_8l_8h_big_cond_dr03_dan():
  """"""
  hparams = imagetransformerpp_sep_channels_8l_8h()
  hparams.hidden_size = 512
  hparams.num_heads = 8
  hparams.filter_size = 2048
  hparams.batch_size = 4
  hparams.max_length = 3075
  hparams.layer_prepostprocess_dropout = 0.3
  hparams.layer_preprocess_sequence = "none"
  hparams.layer_postprocess_sequence = "dan"
  hparams.summarize_grads = True
  hparams.learning_rate = 0.01
  return hparams
1037130
def rnn(name, input, state, kernel, bias, new_state, number_of_gates = 2):
    ''' 
    '''

    nn = Build(name)
    nn.tanh(
        nn.mad(kernel=kernel, bias=bias,
            x=nn.concat(input, state)),
        out=new_state);
    return nn.layers;
293976
def mtf_image_transformer_base():
  """"""
  hparams = common_hparams.basic_params1()
  hparams.no_data_parallelism = True
  hparams.use_fixed_batch_size = True
  hparams.batch_size = 1
  hparams.max_length = 3072
  hparams.hidden_size = 256
  hparams.label_smoothing = 0.0
  # 8-way model-parallelism
  hparams.add_hparam("mesh_shape", "batch:8")
  hparams.add_hparam("layout", "batch:batch")
  hparams.add_hparam("mtf_mode", True)
  hparams.add_hparam("num_heads", 8)
  hparams.add_hparam("filter_size", 1024)
  hparams.add_hparam("num_encoder_layers", 0)
  hparams.add_hparam("num_decoder_layers", 6)
  hparams.add_hparam("attention_key_size", 256)
  hparams.add_hparam("attention_value_size", 256)
  # Share weights between input and target embeddings
  hparams.shared_embedding = True

  # mixture of experts hparams
  hparams.add_hparam("ffn_layer", "dense_relu_dense")
  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

  hparams.shared_embedding_and_softmax_weights = True
  hparams.optimizer = "Adafactor"
  hparams.learning_rate_schedule = "rsqrt_decay"
  hparams.learning_rate_warmup_steps = 10000
  hparams.add_hparam("d_kv", 64)
  hparams.add_hparam("d_ff", 2048)

  # Image related hparams
  hparams.add_hparam("img_len", 32)
  hparams.add_hparam("num_channels", 3)
  hparams.add_hparam("unconditional", True)

  # Local Attention related params
  hparams.add_hparam("block_length", 128)
  hparams.add_hparam("block_height", 16)
  hparams.add_hparam("block_width", 16)
  hparams.add_hparam("attention_type", "local1d")
  return hparams
1171227
def construct_model(images,
                    actions=None,
                    context_frames=2,
                    hparams=None,
                    is_training=True):
  """"""

  pred_depth = 20

  enc_out_all, pred_out_all, van_out_all, van_on_enc_all = [], [], [], []

  lstm_states = [None] * (pred_depth + 2)

  enc_out = encoder_vgg(
      images[0], hparams.enc_size, False, scope_prefix='timestep/',
      hparams=hparams, is_training=is_training)
  enc_out = tf.identity(enc_out, 'enc_out')
  enc_out_all.append(enc_out)

  num_timesteps = len(actions) - 1
  sum_freq = int(num_timesteps / 4 + 1)

  reuse = False
  for timestep, action in zip(range(len(actions) - 1), actions[:-1]):
    done_warm_start = timestep > context_frames - 1

    with tf.variable_scope('timestep', reuse=reuse):
      if done_warm_start:
        pred_input = pred_out_all[-1]
      else:
        pred_input = enc_out_all[-1]
      pred_out = predictor(
          pred_input, action, lstm_states, pred_depth, False, hparams=hparams)
      pred_out = tf.identity(pred_out, 'pred_out')
      if timestep % sum_freq == 0:  # and not hparams.use_tpu:
        tf.summary.histogram('pred_out', pred_out)
      pred_out_all.append(pred_out)

      if timestep % sum_freq == 0:  # and not hparams.use_tpu:
        tf.summary.histogram('lstm_state', lstm_states[0])
      van_out, _, _ = van(
          enc_out_all[0],
          images[0],
          pred_out,
          images[timestep + 1],
          tf.AUTO_REUSE,
          hparams=hparams)
      van_out = tf.identity(van_out, 'van_out')
      van_out_all.append(van_out)

      enc_out = encoder_vgg(
          images[timestep + 1], hparams.enc_size, True, hparams=hparams,
          is_training=is_training)
      enc_out = tf.identity(enc_out, 'enc_out')
      if timestep % sum_freq == 0:  # and not hparams.use_tpu:
        tf.summary.histogram('enc_out', enc_out)
      enc_out_all.append(enc_out)

      van_input = images[0]
      enc_noise = tf.zeros_like(enc_out)
      if timestep % sum_freq == 0:  # and not hparams.use_tpu:
        tf.summary.histogram('enc_noise', enc_noise)
      van_on_enc, _, _ = van(
          enc_out_all[0],
          van_input,
          enc_out + enc_noise,
          images[timestep + 1],
          tf.AUTO_REUSE,
          hparams=hparams)
      van_on_enc = tf.identity(van_on_enc, 'van_on_enc')
      van_on_enc_all.append(van_on_enc)

      reuse = True

  return enc_out_all, pred_out_all, van_out_all, van_on_enc_all
85194
def imagetransformer1d_base_8l_64by64():
  """"""
  hparams = image_transformer_base()
  hparams.num_heads = 8
  hparams.hidden_size = 512
  hparams.filter_size = 2048
  hparams.num_decoder_layers = 8
  hparams.batch_size = 1
  hparams.block_length = 512
  hparams.block_width = 768
  hparams.layer_prepostprocess_dropout = 0.1
  hparams.max_length = 14000
  hparams.unconditional = int(False)
  return hparams