easy-cs-rec-custommodel 0.8.6__py2.py3-none-any.whl

easy_rec/python/layers/keras/attention.py

@@ -0,0 +1,267 @@
+# -*- encoding:utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+"""Attention layers that can be used in sequence DNN/CNN models.
+
+This file follows the terminology of https://arxiv.org/abs/1706.03762 Figure 2.
+Attention is formed by three tensors: Query, Key and Value.
+"""
+import tensorflow as tf
+from tensorflow.python.keras.layers import Layer
+
+
+class Attention(Layer):
+  """Dot-product attention layer, a.k.a. Luong-style attention.
+
+  Inputs are a list with 2 or 3 elements:
+  1. A `query` tensor of shape `(batch_size, Tq, dim)`.
+  2. A `value` tensor of shape `(batch_size, Tv, dim)`.
+  3. A optional `key` tensor of shape `(batch_size, Tv, dim)`. If none
+      supplied, `value` will be used as a `key`.
+
+  The calculation follows the steps:
+  1. Calculate attention scores using `query` and `key` with shape
+      `(batch_size, Tq, Tv)`.
+  2. Use scores to calculate a softmax distribution with shape
+      `(batch_size, Tq, Tv)`.
+  3. Use the softmax distribution to create a linear combination of `value`
+      with shape `(batch_size, Tq, dim)`.
+
+  Args:
+      use_scale: If `True`, will create a scalar variable to scale the
+          attention scores.
+      dropout: Float between 0 and 1. Fraction of the units to drop for the
+          attention scores. Defaults to `0.0`.
+      seed: A Python integer to use as random seed in case of `dropout`.
+      score_mode: Function to use to compute attention scores, one of
+          `{"dot", "concat"}`. `"dot"` refers to the dot product between the
+          query and key vectors. `"concat"` refers to the hyperbolic tangent
+          of the concatenation of the `query` and `key` vectors.
+
+  Call Args:
+      inputs: List of the following tensors:
+          - `query`: Query tensor of shape `(batch_size, Tq, dim)`.
+          - `value`: Value tensor of shape `(batch_size, Tv, dim)`.
+          - `key`: Optional key tensor of shape `(batch_size, Tv, dim)`. If
+              not given, will use `value` for both `key` and `value`, which is
+              the most common case.
+      mask: List of the following tensors:
+          - `query_mask`: A boolean mask tensor of shape `(batch_size, Tq)`.
+              If given, the output will be zero at the positions where
+              `mask==False`.
+          - `value_mask`: A boolean mask tensor of shape `(batch_size, Tv)`.
+              If given, will apply the mask such that values at positions
+               where `mask==False` do not contribute to the result.
+      return_attention_scores: bool, it `True`, returns the attention scores
+          (after masking and softmax) as an additional output argument.
+      training: Python boolean indicating whether the layer should behave in
+          training mode (adding dropout) or in inference mode (no dropout).
+      use_causal_mask: Boolean. Set to `True` for decoder self-attention. Adds
+          a mask such that position `i` cannot attend to positions `j > i`.
+          This prevents the flow of information from the future towards the
+          past. Defaults to `False`.
+
+  Output:
+      Attention outputs of shape `(batch_size, Tq, dim)`.
+      (Optional) Attention scores after masking and softmax with shape
+          `(batch_size, Tq, Tv)`.
+  """
+
+  def __init__(self, params, name='attention', reuse=None, **kwargs):
+    super(Attention, self).__init__(name=name, **kwargs)
+    self.use_scale = params.get_or_default('use_scale', False)
+    self.scale_by_dim = params.get_or_default('scale_by_dim', False)
+    self.score_mode = params.get_or_default('score_mode', 'dot')
+    if self.score_mode not in ['dot', 'concat']:
+      raise ValueError('Invalid value for argument score_mode. '
+                       "Expected one of {'dot', 'concat'}. "
+                       'Received: score_mode=%s' % self.score_mode)
+    self.dropout = params.get_or_default('dropout', 0.0)
+    self.seed = params.get_or_default('seed', None)
+    self.scale = None
+    self.concat_score_weight = None
+    self._return_attention_scores = params.get_or_default(
+        'return_attention_scores', False)
+    self.use_causal_mask = params.get_or_default('use_causal_mask', False)
+
+  @property
+  def return_attention_scores(self):
+    return self._return_attention_scores
+
+  def build(self, input_shape):
+    self._validate_inputs(input_shape)
+    if self.use_scale:
+      self.scale = self.add_weight(
+          name='scale',
+          shape=(),
+          initializer='ones',
+          dtype=self.dtype,
+          trainable=True,
+      )
+    if self.score_mode == 'concat':
+      self.concat_score_weight = self.add_weight(
+          name='concat_score_weight',
+          shape=(),
+          initializer='ones',
+          dtype=self.dtype,
+          trainable=True,
+      )
+    super(Attention, self).build(input_shape)  # Be sure to call this somewhere!
+
+  def _calculate_scores(self, query, key):
+    """Calculates attention scores as a query-key dot product.
+
+    Args:
+        query: Query tensor of shape `(batch_size, Tq, dim)`.
+        key: Key tensor of shape `(batch_size, Tv, dim)`.
+
+    Returns:
+        Tensor of shape `(batch_size, Tq, Tv)`.
+    """
+    if self.score_mode == 'dot':
+      scores = tf.matmul(query, tf.transpose(key, [0, 2, 1]))
+      if self.scale is not None:
+        scores *= self.scale
+      elif self.scale_by_dim:
+        dk = tf.cast(tf.shape(key)[-1], tf.float32)
+        scores /= tf.math.sqrt(dk)
+    elif self.score_mode == 'concat':
+      # Reshape tensors to enable broadcasting.
+      # Reshape into [batch_size, Tq, 1, dim].
+      q_reshaped = tf.expand_dims(query, axis=-2)
+      # Reshape into [batch_size, 1, Tv, dim].
+      k_reshaped = tf.expand_dims(key, axis=-3)
+      if self.scale is not None:
+        scores = self.concat_score_weight * tf.reduce_sum(
+            tf.tanh(self.scale * (q_reshaped + k_reshaped)), axis=-1)
+      else:
+        scores = self.concat_score_weight * tf.reduce_sum(
+            tf.tanh(q_reshaped + k_reshaped), axis=-1)
+    return scores
+
+  def _apply_scores(self, scores, value, scores_mask=None, training=False):
+    """Applies attention scores to the given value tensor.
+
+    To use this method in your attention layer, follow the steps:
+
+    * Use `query` tensor of shape `(batch_size, Tq)` and `key` tensor of
+        shape `(batch_size, Tv)` to calculate the attention `scores`.
+    * Pass `scores` and `value` tensors to this method. The method applies
+        `scores_mask`, calculates
+        `attention_distribution = softmax(scores)`, then returns
+        `matmul(attention_distribution, value).
+    * Apply `query_mask` and return the result.
+
+    Args:
+        scores: Scores float tensor of shape `(batch_size, Tq, Tv)`.
+        value: Value tensor of shape `(batch_size, Tv, dim)`.
+        scores_mask: A boolean mask tensor of shape `(batch_size, 1, Tv)`
+            or `(batch_size, Tq, Tv)`. If given, scores at positions where
+            `scores_mask==False` do not contribute to the result. It must
+            contain at least one `True` value in each line along the last
+            dimension.
+        training: Python boolean indicating whether the layer should behave
+            in training mode (adding dropout) or in inference mode
+            (no dropout).
+
+    Returns:
+        Tensor of shape `(batch_size, Tq, dim)`.
+        Attention scores after masking and softmax with shape
+            `(batch_size, Tq, Tv)`.
+    """
+    if scores_mask is not None:
+      padding_mask = tf.logical_not(scores_mask)
+      # Bias so padding positions do not contribute to attention
+      # distribution.  Note 65504. is the max float16 value.
+      max_value = 65504.0 if scores.dtype == 'float16' else 1.0e9
+      scores -= max_value * tf.cast(padding_mask, dtype=scores.dtype)
+
+    weights = tf.nn.softmax(scores, axis=-1)
+    if training and self.dropout > 0:
+      weights = tf.nn.dropout(weights, 1.0 - self.dropout, seed=self.seed)
+    return tf.matmul(weights, value), weights
+
+  def _calculate_score_mask(self, scores, v_mask, use_causal_mask):
+    if use_causal_mask:
+      # Creates a lower triangular mask, so position i cannot attend to
+      # positions j > i. This prevents the flow of information from the
+      # future into the past.
+      score_shape = tf.shape(scores)
+      # causal_mask_shape = [1, Tq, Tv].
+      mask_shape = (1, score_shape[-2], score_shape[-1])
+      ones_mask = tf.ones(shape=mask_shape, dtype='int32')
+      row_index = tf.cumsum(ones_mask, axis=-2)
+      col_index = tf.cumsum(ones_mask, axis=-1)
+      causal_mask = tf.greater_equal(row_index, col_index)
+
+      if v_mask is not None:
+        # Mask of shape [batch_size, 1, Tv].
+        v_mask = tf.expand_dims(v_mask, axis=-2)
+        return tf.logical_and(v_mask, causal_mask)
+      return causal_mask
+    else:
+      # If not using causal mask, return the value mask as is,
+      # or None if the value mask is not provided.
+      return v_mask
+
+  def call(self, inputs, mask=None, training=False, **kwargs):
+    self._validate_inputs(inputs=inputs, mask=mask)
+    q = inputs[0]
+    v = inputs[1]
+    k = inputs[2] if len(inputs) > 2 else v
+    q_mask = mask[0] if mask else None
+    v_mask = mask[1] if mask else None
+    scores = self._calculate_scores(query=q, key=k)
+    scores_mask = self._calculate_score_mask(scores, v_mask,
+                                             self.use_causal_mask)
+    result, attention_scores = self._apply_scores(
+        scores=scores, value=v, scores_mask=scores_mask, training=training)
+    if q_mask is not None:
+      # Mask of shape [batch_size, Tq, 1].
+      q_mask = tf.expand_dims(q_mask, axis=-1)
+      result *= tf.cast(q_mask, dtype=result.dtype)
+    if self._return_attention_scores:
+      return result, attention_scores
+    return result
+
+  def compute_mask(self, inputs, mask=None):
+    self._validate_inputs(inputs=inputs, mask=mask)
+    if mask is None or mask[0] is None:
+      return None
+    return tf.convert_to_tensor(mask[0])
+
+  def compute_output_shape(self, input_shape):
+    """Returns shape of value tensor dim, but for query tensor length."""
+    return list(input_shape[0][:-1]), input_shape[1][-1]
+
+  def _validate_inputs(self, inputs, mask=None):
+    """Validates arguments of the call method."""
+    class_name = self.__class__.__name__
+    if not isinstance(inputs, list):
+      raise ValueError('{class_name} layer must be called on a list of inputs, '
+                       'namely [query, value] or [query, value, key]. '
+                       'Received: inputs={inputs}.'.format(
+                           class_name=class_name, inputs=inputs))
+    if len(inputs) < 2 or len(inputs) > 3:
+      raise ValueError('%s layer accepts inputs list of length 2 or 3, '
+                       'namely [query, value] or [query, value, key]. '
+                       'Received length: %d.' % (class_name, len(inputs)))
+    if mask is not None:
+      if not isinstance(mask, list):
+        raise ValueError(
+            '{class_name} layer mask must be a list, '
+            'namely [query_mask, value_mask]. Received: mask={mask}.'.format(
+                class_name=class_name, mask=mask))
+      if len(mask) < 2 or len(mask) > 3:
+        raise ValueError(
+            '{class_name} layer accepts mask list of length 2 or 3. '
+            'Received: inputs={inputs}, mask={mask}.'.format(
+                class_name=class_name, inputs=inputs, mask=mask))
+
+  def get_config(self):
+    base_config = super(Attention, self).get_config()
+    config = {
+        'use_scale': self.use_scale,
+        'score_mode': self.score_mode,
+        'dropout': self.dropout,
+    }
+    return dict(list(base_config.items()) + list(config.items()))

easy_rec/python/layers/keras/auxiliary_loss.py

@@ -0,0 +1,47 @@
+# -*- encoding:utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import logging
+
+import tensorflow as tf
+
+from easy_rec.python.loss import contrastive_loss
+
+
+class AuxiliaryLoss(tf.keras.layers.Layer):
+  """Compute auxiliary loss, usually use for contrastive learning."""
+
+  def __init__(self, params, name='auxiliary_loss', reuse=None, **kwargs):
+    super(AuxiliaryLoss, self).__init__(name=name, **kwargs)
+    params.check_required('loss_type')
+    self.loss_type = params.get_or_default('loss_type', None)
+    self.loss_weight = params.get_or_default('loss_weight', 1.0)
+    logging.info('init layer `%s` with loss type: %s and weight: %f' %
+                 (self.name, self.loss_type, self.loss_weight))
+    self.temperature = params.get_or_default('temperature', 0.1)
+
+  def call(self, inputs, training=None, **kwargs):
+    if self.loss_type is None:
+      logging.warning('loss_type is None in auxiliary loss layer')
+      return 0
+
+    loss_dict = kwargs['loss_dict']
+    loss_value = 0
+
+    if self.loss_type == 'l2_loss':
+      x1, x2 = inputs
+      loss = contrastive_loss.l2_loss(x1, x2)
+      loss_value = loss if self.loss_weight == 1.0 else loss * self.loss_weight
+      loss_dict['%s_l2_loss' % self.name] = loss_value
+    elif self.loss_type == 'info_nce':
+      query, positive = inputs
+      loss = contrastive_loss.info_nce_loss(
+          query, positive, temperature=self.temperature)
+      loss_value = loss if self.loss_weight == 1.0 else loss * self.loss_weight
+      loss_dict['%s_info_nce_loss' % self.name] = loss_value
+    elif self.loss_type == 'nce_loss':
+      x1, x2 = inputs
+      loss = contrastive_loss.nce_loss(x1, x2, temperature=self.temperature)
+      loss_value = loss if self.loss_weight == 1.0 else loss * self.loss_weight
+      loss_dict['%s_nce_loss' % self.name] = loss_value
+
+    return loss_value

easy_rec/python/layers/keras/blocks.py

@@ -0,0 +1,262 @@
+# -*- encoding:utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+"""Convenience blocks for building models."""
+import logging
+
+import tensorflow as tf
+from tensorflow.python.keras.initializers import Constant
+from tensorflow.python.keras.layers import Dense
+from tensorflow.python.keras.layers import Dropout
+from tensorflow.python.keras.layers import Lambda
+from tensorflow.python.keras.layers import Layer
+
+from easy_rec.python.layers.keras.activation import activation_layer
+from easy_rec.python.layers.utils import Parameter
+from easy_rec.python.utils.shape_utils import pad_or_truncate_sequence
+from easy_rec.python.utils.tf_utils import add_elements_to_collection
+
+if tf.__version__ >= '2.0':
+  tf = tf.compat.v1
+
+
+class MLP(Layer):
+  """Sequential multi-layer perceptron (MLP) block.
+
+  Attributes:
+    units: Sequential list of layer sizes.
+    use_bias: Whether to include a bias term.
+    activation: Type of activation to use on all except the last layer.
+    final_activation: Type of activation to use on last layer.
+    **kwargs: Extra args passed to the Keras Layer base class.
+  """
+
+  def __init__(self, params, name='mlp', reuse=None, **kwargs):
+    super(MLP, self).__init__(name=name, **kwargs)
+    self.layer_name = name  # for add to output
+    params.check_required('hidden_units')
+    use_bn = params.get_or_default('use_bn', True)
+    use_final_bn = params.get_or_default('use_final_bn', True)
+    use_bias = params.get_or_default('use_bias', False)
+    use_final_bias = params.get_or_default('use_final_bias', False)
+    dropout_rate = list(params.get_or_default('dropout_ratio', []))
+    activation = params.get_or_default('activation', 'relu')
+    initializer = params.get_or_default('initializer', 'he_uniform')
+    final_activation = params.get_or_default('final_activation', None)
+    use_bn_after_act = params.get_or_default('use_bn_after_activation', False)
+    units = list(params.hidden_units)
+    logging.info(
+        'MLP(%s) units: %s, dropout: %r, activate=%s, use_bn=%r, final_bn=%r,'
+        ' final_activate=%s, bias=%r, initializer=%s, bn_after_activation=%r' %
+        (name, units, dropout_rate, activation, use_bn, use_final_bn,
+         final_activation, use_bias, initializer, use_bn_after_act))
+    assert len(units) > 0, 'MLP(%s) takes at least one hidden units' % name
+    self.reuse = reuse
+    self.add_to_outputs = params.get_or_default('add_to_outputs', False)
+
+    num_dropout = len(dropout_rate)
+    self._sub_layers = []
+    for i, num_units in enumerate(units[:-1]):
+      name = 'layer_%d' % i
+      drop_rate = dropout_rate[i] if i < num_dropout else 0.0
+      self.add_rich_layer(num_units, use_bn, drop_rate, activation, initializer,
+                          use_bias, use_bn_after_act, name,
+                          params.l2_regularizer)
+
+    n = len(units) - 1
+    drop_rate = dropout_rate[n] if num_dropout > n else 0.0
+    name = 'layer_%d' % n
+    self.add_rich_layer(units[-1], use_final_bn, drop_rate, final_activation,
+                        initializer, use_final_bias, use_bn_after_act, name,
+                        params.l2_regularizer)
+
+  def add_rich_layer(self,
+                     num_units,
+                     use_bn,
+                     dropout_rate,
+                     activation,
+                     initializer,
+                     use_bias,
+                     use_bn_after_activation,
+                     name,
+                     l2_reg=None):
+    act_layer = activation_layer(activation, name='%s/act' % name)
+    if use_bn and not use_bn_after_activation:
+      dense = Dense(
+          units=num_units,
+          use_bias=use_bias,
+          kernel_initializer=initializer,
+          kernel_regularizer=l2_reg,
+          name='%s/dense' % name)
+      self._sub_layers.append(dense)
+      bn = tf.keras.layers.BatchNormalization(
+          name='%s/bn' % name, trainable=True)
+      self._sub_layers.append(bn)
+      self._sub_layers.append(act_layer)
+    else:
+      dense = Dense(
+          num_units,
+          use_bias=use_bias,
+          kernel_initializer=initializer,
+          kernel_regularizer=l2_reg,
+          name='%s/dense' % name)
+      self._sub_layers.append(dense)
+      self._sub_layers.append(act_layer)
+      if use_bn and use_bn_after_activation:
+        bn = tf.keras.layers.BatchNormalization(name='%s/bn' % name)
+        self._sub_layers.append(bn)
+
+    if 0.0 < dropout_rate < 1.0:
+      dropout = Dropout(dropout_rate, name='%s/dropout' % name)
+      self._sub_layers.append(dropout)
+    elif dropout_rate >= 1.0:
+      raise ValueError('invalid dropout_ratio: %.3f' % dropout_rate)
+
+  def call(self, x, training=None, **kwargs):
+    """Performs the forward computation of the block."""
+    for layer in self._sub_layers:
+      cls = layer.__class__.__name__
+      if cls in ('Dropout', 'BatchNormalization', 'Dice'):
+        x = layer(x, training=training)
+        if cls in ('BatchNormalization', 'Dice') and training:
+          add_elements_to_collection(layer.updates, tf.GraphKeys.UPDATE_OPS)
+      else:
+        x = layer(x)
+    if self.add_to_outputs and 'prediction_dict' in kwargs:
+      outputs = kwargs['prediction_dict']
+      outputs[self.layer_name] = tf.squeeze(x, axis=1)
+      logging.info('add `%s` to model outputs' % self.layer_name)
+    return x
+
+
+class Highway(Layer):
+
+  def __init__(self, params, name='highway', reuse=None, **kwargs):
+    super(Highway, self).__init__(name=name, **kwargs)
+    self.emb_size = params.get_or_default('emb_size', None)
+    self.num_layers = params.get_or_default('num_layers', 1)
+    self.activation = params.get_or_default('activation', 'relu')
+    self.dropout_rate = params.get_or_default('dropout_rate', 0.0)
+    self.init_gate_bias = params.get_or_default('init_gate_bias', -3.0)
+    self.act_layer = activation_layer(self.activation)
+    self.dropout_layer = Dropout(
+        self.dropout_rate) if self.dropout_rate > 0.0 else None
+    self.project_layer = None
+    self.gate_bias_initializer = Constant(self.init_gate_bias)
+    self.gates = []  # T
+    self.transforms = []  # H
+    self.multiply_layer = tf.keras.layers.Multiply()
+    self.add_layer = tf.keras.layers.Add()
+
+  def build(self, input_shape):
+    dim = input_shape[-1]
+    if self.emb_size is not None and dim != self.emb_size:
+      self.project_layer = Dense(self.emb_size, name='input_projection')
+      dim = self.emb_size
+    self.carry_gate = Lambda(lambda x: 1.0 - x, output_shape=(dim,))
+    for i in range(self.num_layers):
+      gate = Dense(
+          units=dim,
+          bias_initializer=self.gate_bias_initializer,
+          activation='sigmoid',
+          name='gate_%d' % i)
+      self.gates.append(gate)
+      self.transforms.append(Dense(units=dim))
+
+  def call(self, inputs, training=None, **kwargs):
+    value = inputs
+    if self.project_layer is not None:
+      value = self.project_layer(inputs)
+    for i in range(self.num_layers):
+      gate = self.gates[i](value)
+      transformed = self.act_layer(self.transforms[i](value))
+      if self.dropout_layer is not None:
+        transformed = self.dropout_layer(transformed, training=training)
+      transformed_gated = self.multiply_layer([gate, transformed])
+      identity_gated = self.multiply_layer([self.carry_gate(gate), value])
+      value = self.add_layer([transformed_gated, identity_gated])
+    return value
+
+
+class Gate(Layer):
+  """Weighted sum gate."""
+
+  def __init__(self, params, name='gate', reuse=None, **kwargs):
+    super(Gate, self).__init__(name=name, **kwargs)
+    self.weight_index = params.get_or_default('weight_index', 0)
+    if params.has_field('mlp'):
+      mlp_cfg = Parameter.make_from_pb(params.mlp)
+      mlp_cfg.l2_regularizer = params.l2_regularizer
+      self.top_mlp = MLP(mlp_cfg, name='top_mlp')
+    else:
+      self.top_mlp = None
+
+  def call(self, inputs, training=None, **kwargs):
+    assert len(
+        inputs
+    ) > 1, 'input of Gate layer must be a list containing at least 2 elements'
+    weights = inputs[self.weight_index]
+    j = 0
+    for i, x in enumerate(inputs):
+      if i == self.weight_index:
+        continue
+      if j == 0:
+        output = weights[:, j, None] * x
+      else:
+        output += weights[:, j, None] * x
+      j += 1
+    if self.top_mlp is not None:
+      output = self.top_mlp(output, training=training)
+    return output
+
+
+class TextCNN(Layer):
+  """Text CNN Model.
+
+  References
+  - [Convolutional Neural Networks for Sentence Classification](https://arxiv.org/abs/1408.5882)
+  """
+
+  def __init__(self, params, name='text_cnn', reuse=None, **kwargs):
+    super(TextCNN, self).__init__(name=name, **kwargs)
+    self.config = params.get_pb_config()
+    self.pad_seq_length = self.config.pad_sequence_length
+    if self.pad_seq_length <= 0:
+      logging.warning(
+          'run text cnn with pad_sequence_length <= 0, the predict of model may be unstable'
+      )
+    self.conv_layers = []
+    self.pool_layer = tf.keras.layers.GlobalMaxPool1D()
+    self.concat_layer = tf.keras.layers.Concatenate(axis=-1)
+    for size, filters in zip(self.config.filter_sizes, self.config.num_filters):
+      conv = tf.keras.layers.Conv1D(
+          filters=int(filters),
+          kernel_size=int(size),
+          activation=self.config.activation)
+      self.conv_layers.append(conv)
+    if self.config.HasField('mlp'):
+      p = Parameter.make_from_pb(self.config.mlp)
+      p.l2_regularizer = params.l2_regularizer
+      self.mlp = MLP(p, name='mlp', reuse=reuse)
+    else:
+      self.mlp = None
+
+  def call(self, inputs, training=None, **kwargs):
+    """Input shape: 3D tensor with shape: `(batch_size, steps, input_dim)."""
+    assert isinstance(inputs, (list, tuple))
+    assert len(inputs) >= 2
+    seq_emb, seq_len = inputs[:2]
+
+    if self.pad_seq_length > 0:
+      seq_emb, seq_len = pad_or_truncate_sequence(seq_emb, seq_len,
+                                                  self.pad_seq_length)
+    pooled_outputs = []
+    for layer in self.conv_layers:
+      conv = layer(seq_emb)
+      pooled = self.pool_layer(conv)
+      pooled_outputs.append(pooled)
+    net = self.concat_layer(pooled_outputs)
+    if self.mlp is not None:
+      output = self.mlp(net, training=training)
+    else:
+      output = net
+    return output