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

easy_rec/python/layers/keras/layer_norm.py

@@ -0,0 +1,364 @@
+"""Layer Normalization layer."""
+import tensorflow as tf
+from tensorflow.python.keras import constraints
+from tensorflow.python.keras import initializers
+from tensorflow.python.keras import regularizers
+from tensorflow.python.keras.layers import Layer
+
+
+def validate_axis(axis, input_shape):
+  """Validate an axis value and returns its standardized form.
+
+  Args:
+    axis: Value to validate. Can be an integer or a list/tuple of integers.
+      Integers may be negative.
+    input_shape: Reference input shape that the axis/axes refer to.
+
+  Returns:
+    Normalized form of `axis`, i.e. a list with all-positive values.
+  """
+  input_shape = tf.TensorShape(input_shape)
+  rank = input_shape.ndims
+  if not rank:
+    raise ValueError(
+        'Input has undefined rank. Received: input_shape={input_shape}'.format(
+            input_shape=input_shape))
+
+  # Convert axis to list and resolve negatives
+  if isinstance(axis, int):
+    axis = [axis]
+  else:
+    axis = list(axis)
+  for idx, x in enumerate(axis):
+    if x < 0:
+      axis[idx] = rank + x
+
+  # Validate axes
+  for x in axis:
+    if x < 0 or x >= rank:
+      raise ValueError('Invalid value for `axis` argument. '
+                       'Expected 0 <= axis < inputs.rank (with '
+                       'inputs.rank={rank}). Received: axis={axis}'.format(
+                           rank=rank, axis=tuple(axis)))
+  if len(axis) != len(set(axis)):
+    raise ValueError('Duplicate axis: {axis}'.format(axis=tuple(axis)))
+  return axis
+
+
+class LayerNormalization(Layer):
+  """Layer normalization layer (Ba et al., 2016).
+
+  Normalize the activations of the previous layer for each given example in a
+  batch independently, rather than across a batch like Batch Normalization.
+  i.e. applies a transformation that maintains the mean activation within each
+  example close to 0 and the activation standard deviation close to 1.
+
+  Given a tensor `inputs`, moments are calculated and normalization
+  is performed across the axes specified in `axis`.
+
+  Example:
+  >>> data = tf.constant(np.arange(10).reshape(5, 2) * 10, dtype=tf.float32)
+  >>> print(data)
+  tf.Tensor(
+  [[ 0. 10.]
+   [20. 30.]
+   [40. 50.]
+   [60. 70.]
+   [80. 90.]], shape=(5, 2), dtype=float32)
+
+  >>> layer = tf.keras.layers.LayerNormalization(axis=1)
+  >>> output = layer(data)
+  >>> print(output)
+  tf.Tensor(
+  [[-1. 1.]
+   [-1. 1.]
+   [-1. 1.]
+   [-1. 1.]
+   [-1. 1.]], shape=(5, 2), dtype=float32)
+
+  Notice that with Layer Normalization the normalization happens across the
+  axes *within* each example, rather than across different examples in the
+  batch.
+
+  If `scale` or `center` are enabled, the layer will scale the normalized
+  outputs by broadcasting them with a trainable variable `gamma`, and center
+  the outputs by broadcasting with a trainable variable `beta`. `gamma` will
+  default to a ones tensor and `beta` will default to a zeros tensor, so that
+  centering and scaling are no-ops before training has begun.
+
+  So, with scaling and centering enabled the normalization equations
+  are as follows:
+
+  Let the intermediate activations for a mini-batch to be the `inputs`.
+
+  For each sample `x_i` in `inputs` with `k` features, we compute the mean and
+  variance of the sample:
+
+  ```python
+  mean_i = sum(x_i[j] for j in range(k)) / k
+  var_i = sum((x_i[j] - mean_i) ** 2 for j in range(k)) / k
+  ```
+
+  and then compute a normalized `x_i_normalized`, including a small factor
+  `epsilon` for numerical stability.
+
+  ```python
+  x_i_normalized = (x_i - mean_i) / sqrt(var_i + epsilon)
+  ```
+
+  And finally `x_i_normalized ` is linearly transformed by `gamma` and `beta`,
+  which are learned parameters:
+
+  ```python
+  output_i = x_i_normalized * gamma + beta
+  ```
+
+  `gamma` and `beta` will span the axes of `inputs` specified in `axis`, and
+  this part of the inputs' shape must be fully defined.
+
+  For example:
+  >>> layer = tf.keras.layers.LayerNormalization(axis=[1, 2, 3])
+  >>> layer.build([5, 20, 30, 40])
+  >>> print(layer.beta.shape)
+  (20, 30, 40)
+  >>> print(layer.gamma.shape)
+  (20, 30, 40)
+
+  Note that other implementations of layer normalization may choose to define
+  `gamma` and `beta` over a separate set of axes from the axes being
+  normalized across. For example, Group Normalization
+  ([Wu et al. 2018](https://arxiv.org/abs/1803.08494)) with group size of 1
+  corresponds to a Layer Normalization that normalizes across height, width,
+  and channel and has `gamma` and `beta` span only the channel dimension.
+  So, this Layer Normalization implementation will not match a Group
+  Normalization layer with group size set to 1.
+
+  Args:
+    axis: Integer or List/Tuple. The axis or axes to normalize across.
+      Typically, this is the features axis/axes. The left-out axes are
+      typically the batch axis/axes. `-1` is the last dimension in the
+      input. Defaults to `-1`.
+    epsilon: Small float added to variance to avoid dividing by zero. Defaults
+      to 1e-3
+    center: If True, add offset of `beta` to normalized tensor. If False,
+      `beta` is ignored. Defaults to `True`.
+    scale: If True, multiply by `gamma`. If False, `gamma` is not used.
+      When the next layer is linear (also e.g. `nn.relu`), this can be
+      disabled since the scaling will be done by the next layer.
+      Defaults to `True`.
+    beta_initializer: Initializer for the beta weight. Defaults to zeros.
+    gamma_initializer: Initializer for the gamma weight. Defaults to ones.
+    beta_regularizer: Optional regularizer for the beta weight. None by
+      default.
+    gamma_regularizer: Optional regularizer for the gamma weight. None by
+      default.
+    beta_constraint: Optional constraint for the beta weight. None by default.
+    gamma_constraint: Optional constraint for the gamma weight. None by
+      default.
+
+  Input shape:
+    Arbitrary. Use the keyword argument `input_shape` (tuple of
+    integers, does not include the samples axis) when using this layer as the
+    first layer in a model.
+
+  Output shape:
+    Same shape as input.
+
+  Reference:
+    - [Lei Ba et al., 2016](https://arxiv.org/abs/1607.06450).
+  """
+
+  def __init__(self,
+               axis=-1,
+               epsilon=1e-3,
+               center=True,
+               scale=True,
+               beta_initializer='zeros',
+               gamma_initializer='ones',
+               beta_regularizer=None,
+               gamma_regularizer=None,
+               beta_constraint=None,
+               gamma_constraint=None,
+               **kwargs):
+    super(LayerNormalization, self).__init__(**kwargs)
+    if isinstance(axis, (list, tuple)):
+      self.axis = list(axis)
+    elif isinstance(axis, int):
+      self.axis = axis
+    else:
+      raise TypeError('Expected an int or a list/tuple of ints for the '
+                      "argument 'axis', but received: %r" % axis)
+
+    self.epsilon = epsilon
+    self.center = center
+    self.scale = scale
+    self.beta_initializer = initializers.get(beta_initializer)
+    self.gamma_initializer = initializers.get(gamma_initializer)
+    self.beta_regularizer = regularizers.get(beta_regularizer)
+    self.gamma_regularizer = regularizers.get(gamma_regularizer)
+    self.beta_constraint = constraints.get(beta_constraint)
+    self.gamma_constraint = constraints.get(gamma_constraint)
+
+    self.supports_masking = True
+
+    # Indicates whether a faster fused implementation can be used. This will
+    # be set to True or False in build()"
+    self._fused = None
+
+  def _fused_can_be_used(self, ndims):
+    """Returns false if fused implementation cannot be used.
+
+    Check if the axis is contiguous and can be collapsed into the last axis.
+    The self.axis is assumed to have no duplicates.
+    """
+    if not tf.test.is_gpu_available():
+      return False
+    axis = sorted(self.axis)
+    can_use_fused = False
+
+    if axis[-1] == ndims - 1 and axis[-1] - axis[0] == len(axis) - 1:
+      can_use_fused = True
+
+    # fused_batch_norm will silently raise epsilon to be at least 1.001e-5,
+    # so we cannot used the fused version if epsilon is below that value.
+    # Also, the variable dtype must be float32, as fused_batch_norm only
+    # supports float32 variables.
+    if self.epsilon < 1.001e-5 or self.dtype != 'float32':
+      can_use_fused = False
+
+    return can_use_fused
+
+  def build(self, input_shape):
+    self.axis = validate_axis(self.axis, input_shape)
+    input_shape = tf.TensorShape(input_shape)
+    rank = input_shape.ndims
+
+    param_shape = [input_shape[dim] for dim in self.axis]
+    if self.scale:
+      self.gamma = self.add_weight(
+          name='gamma',
+          shape=param_shape,
+          initializer=self.gamma_initializer,
+          regularizer=self.gamma_regularizer,
+          constraint=self.gamma_constraint,
+          trainable=True,
+      )
+    else:
+      self.gamma = None
+
+    if self.center:
+      self.beta = self.add_weight(
+          name='beta',
+          shape=param_shape,
+          initializer=self.beta_initializer,
+          regularizer=self.beta_regularizer,
+          constraint=self.beta_constraint,
+          trainable=True,
+      )
+    else:
+      self.beta = None
+
+    self._fused = self._fused_can_be_used(rank)
+    super(LayerNormalization,
+          self).build(input_shape)  # Be sure to call this somewhere!
+
+  def call(self, inputs):
+    # Compute the axes along which to reduce the mean / variance
+    input_shape = inputs.shape
+    ndims = len(input_shape)
+
+    # Broadcasting only necessary for norm when the axis is not just
+    # the last dimension
+    broadcast_shape = [1] * ndims
+    for dim in self.axis:
+      broadcast_shape[dim] = input_shape.dims[dim].value
+
+    def _broadcast(v):
+      if (v is not None and len(v.shape) != ndims and self.axis != [ndims - 1]):
+        return tf.reshape(v, broadcast_shape)
+      return v
+
+    if not self._fused:
+      input_dtype = inputs.dtype
+      if (input_dtype in ('float16', 'bfloat16') and self.dtype == 'float32'):
+        # If mixed precision is used, cast inputs to float32 so that
+        # this is at least as numerically stable as the fused version.
+        inputs = tf.cast(inputs, 'float32')
+
+      # Calculate the moments on the last axis (layer activations).
+      mean, variance = tf.nn.moments(inputs, self.axis, keep_dims=True)
+
+      scale, offset = _broadcast(self.gamma), _broadcast(self.beta)
+
+      # Compute layer normalization using the batch_normalization
+      # function.
+      outputs = tf.nn.batch_normalization(
+          inputs,
+          mean,
+          variance,
+          offset=offset,
+          scale=scale,
+          variance_epsilon=self.epsilon,
+      )
+      outputs = tf.cast(outputs, input_dtype)
+    else:
+      # Collapse dims before self.axis, and dims in self.axis
+
+      axis = sorted(self.axis)
+      tensor_shape = tf.shape(inputs)
+      pre_dim = tf.reduce_prod(tensor_shape[:axis[0]])
+      in_dim = tf.reduce_prod(tensor_shape[axis[0]:])
+      squeezed_shape = [1, pre_dim, in_dim, 1]
+      # This fused operation requires reshaped inputs to be NCHW.
+      data_format = 'NCHW'
+
+      inputs = tf.reshape(inputs, squeezed_shape)
+
+      # self.gamma and self.beta have the wrong shape for
+      # fused_batch_norm, so we cannot pass them as the scale and offset
+      # parameters. Therefore, we create two constant tensors in correct
+      # shapes for fused_batch_norm and later construct a separate
+      # calculation on the scale and offset.
+      scale = tf.ones(tf.convert_to_tensor([pre_dim]), dtype=self.dtype)
+      offset = tf.zeros(tf.convert_to_tensor([pre_dim]), dtype=self.dtype)
+
+      # Compute layer normalization using the fused_batch_norm function.
+      outputs, _, _ = tf.compat.v1.nn.fused_batch_norm(
+          inputs,
+          scale=scale,
+          offset=offset,
+          epsilon=self.epsilon,
+          data_format=data_format,
+      )
+
+      outputs = tf.reshape(outputs, tensor_shape)
+
+      scale, offset = _broadcast(self.gamma), _broadcast(self.beta)
+
+      if scale is not None:
+        outputs = outputs * tf.cast(scale, outputs.dtype)
+      if offset is not None:
+        outputs = outputs + tf.cast(offset, outputs.dtype)
+
+    # If some components of the shape got lost due to adjustments, fix that.
+    outputs.set_shape(input_shape)
+    return outputs
+
+  def compute_output_shape(self, input_shape):
+    return input_shape
+
+  def get_config(self):
+    config = {
+        'axis': self.axis,
+        'epsilon': self.epsilon,
+        'center': self.center,
+        'scale': self.scale,
+        'beta_initializer': initializers.serialize(self.beta_initializer),
+        'gamma_initializer': initializers.serialize(self.gamma_initializer),
+        'beta_regularizer': regularizers.serialize(self.beta_regularizer),
+        'gamma_regularizer': regularizers.serialize(self.gamma_regularizer),
+        'beta_constraint': constraints.serialize(self.beta_constraint),
+        'gamma_constraint': constraints.serialize(self.gamma_constraint),
+    }
+    base_config = super(LayerNormalization, self).get_config()
+    return dict(list(base_config.items()) + list(config.items()))

easy_rec/python/layers/keras/mask_net.py

@@ -0,0 +1,166 @@
+# -*- encoding:utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import logging
+
+import tensorflow as tf
+from tensorflow.python.keras.layers import Activation
+from tensorflow.python.keras.layers import Dense
+from tensorflow.python.keras.layers import Layer
+
+from easy_rec.python.layers.keras.blocks import MLP
+from easy_rec.python.layers.keras.layer_norm import LayerNormalization
+from easy_rec.python.layers.utils import Parameter
+
+
+class MaskBlock(Layer):
+  """MaskBlock use in MaskNet.
+
+  Args:
+    projection_dim: project dimension to reduce the computational cost.
+    Default is `None` such that a full (`input_dim` by `aggregation_size`) matrix
+    W is used. If enabled, a low-rank matrix W = U*V will be used, where U
+    is of size `input_dim` by `projection_dim` and V is of size
+    `projection_dim` by `aggregation_size`. `projection_dim` need to be smaller
+    than `aggregation_size`/2 to improve the model efficiency. In practice, we've
+    observed that `projection_dim` = d/4 consistently preserved the
+    accuracy of a full-rank version.
+  """
+
+  def __init__(self, params, name='mask_block', reuse=None, **kwargs):
+    super(MaskBlock, self).__init__(name=name, **kwargs)
+    self.config = params.get_pb_config()
+    self.l2_reg = params.l2_regularizer
+    self._projection_dim = params.get_or_default('projection_dim', None)
+    self.reuse = reuse
+    self.final_relu = Activation('relu', name='relu')
+
+  def build(self, input_shape):
+    if type(input_shape) in (tuple, list):
+      assert len(input_shape) >= 2, 'MaskBlock must has at least two inputs'
+      input_dim = int(input_shape[0][-1])
+      mask_input_dim = int(input_shape[1][-1])
+    else:
+      input_dim, mask_input_dim = input_shape[-1], input_shape[-1]
+    if self.config.HasField('reduction_factor'):
+      aggregation_size = int(mask_input_dim * self.config.reduction_factor)
+    elif self.config.HasField('aggregation_size') is not None:
+      aggregation_size = self.config.aggregation_size
+    else:
+      raise ValueError(
+          'Need one of reduction factor or aggregation size for MaskBlock.')
+
+    self.aggr_layer = Dense(
+        aggregation_size,
+        activation='relu',
+        kernel_initializer='he_uniform',
+        kernel_regularizer=self.l2_reg,
+        name='aggregation')
+    self.weight_layer = Dense(input_dim, name='weights')
+    if self._projection_dim is not None:
+      logging.info('%s project dim is %d', self.name, self._projection_dim)
+      self.project_layer = Dense(
+          self._projection_dim,
+          kernel_regularizer=self.l2_reg,
+          use_bias=False,
+          name='project')
+    if self.config.input_layer_norm:
+      # 推荐在调用MaskBlock之前做好 layer norm，否则每一次调用都需要对input做ln
+      if tf.__version__ >= '2.0':
+        self.input_layer_norm = tf.keras.layers.LayerNormalization(
+            name='input_ln')
+      else:
+        self.input_layer_norm = LayerNormalization(name='input_ln')
+
+    if self.config.HasField('output_size'):
+      self.output_layer = Dense(
+          self.config.output_size, use_bias=False, name='output')
+    if tf.__version__ >= '2.0':
+      self.output_layer_norm = tf.keras.layers.LayerNormalization(
+          name='output_ln')
+    else:
+      self.output_layer_norm = LayerNormalization(name='output_ln')
+    super(MaskBlock, self).build(input_shape)
+
+  def call(self, inputs, training=None, **kwargs):
+    if type(inputs) in (tuple, list):
+      net, mask_input = inputs[:2]
+    else:
+      net, mask_input = inputs, inputs
+
+    if self.config.input_layer_norm:
+      net = self.input_layer_norm(net)
+
+    if self._projection_dim is None:
+      aggr = self.aggr_layer(mask_input)
+    else:
+      u = self.project_layer(mask_input)
+      aggr = self.aggr_layer(u)
+
+    weights = self.weight_layer(aggr)
+    masked_net = net * weights
+
+    if not self.config.HasField('output_size'):
+      return masked_net
+
+    hidden = self.output_layer(masked_net)
+    ln_hidden = self.output_layer_norm(hidden)
+    return self.final_relu(ln_hidden)
+
+
+class MaskNet(Layer):
+  """MaskNet: Introducing Feature-Wise Multiplication to CTR Ranking Models by Instance-Guided Mask.
+
+  Refer: https://arxiv.org/pdf/2102.07619.pdf
+  """
+
+  def __init__(self, params, name='mask_net', reuse=None, **kwargs):
+    super(MaskNet, self).__init__(name=name, **kwargs)
+    self.reuse = reuse
+    self.params = params
+    self.config = params.get_pb_config()
+    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
+
+    self.mask_layers = []
+    for i, block_conf in enumerate(self.config.mask_blocks):
+      params = Parameter.make_from_pb(block_conf)
+      params.l2_regularizer = self.params.l2_regularizer
+      mask_layer = MaskBlock(params, name='block_%d' % i, reuse=self.reuse)
+      self.mask_layers.append(mask_layer)
+
+    if self.config.input_layer_norm:
+      if tf.__version__ >= '2.0':
+        self.input_layer_norm = tf.keras.layers.LayerNormalization(
+            name='input_ln')
+      else:
+        self.input_layer_norm = LayerNormalization(name='input_ln')
+
+  def call(self, inputs, training=None, **kwargs):
+    if self.config.input_layer_norm:
+      inputs = self.input_layer_norm(inputs)
+
+    if self.config.use_parallel:
+      mask_outputs = [
+          mask_layer((inputs, inputs)) for mask_layer in self.mask_layers
+      ]
+      all_mask_outputs = tf.concat(mask_outputs, axis=1)
+      if self.mlp is not None:
+        output = self.mlp(all_mask_outputs, training=training)
+      else:
+        output = all_mask_outputs
+      return output
+    else:
+      net = inputs
+      for i, _ in enumerate(self.config.mask_blocks):
+        mask_layer = self.mask_layers[i]
+        net = mask_layer((net, inputs))
+
+      if self.mlp is not None:
+        output = self.mlp(net, training=training)
+      else:
+        output = net
+      return output