Does creating tensors in the call method of a custom layer create new nodes on the graph each time the custom layer is called?

user_11064604 0
Viewed 104

TLDR

I need random effects at each call of my custom layer. If I create a random mask (and also create a new tensor by shuffling the input data) in the call method instead of the init or build methods, will this create new nodes on the computation graph with every function call?

I am creating a custom CutMix-like layer to use on tabular data. This layer will 1) take a minibatch, 2) create a shuffled version of the original minibatch, 3) replace the values of the original for the shuffled at bernouilli(p). This is often referred to as SwapNoise on Kaggle.

This layer relies on a random mask (drawn bernoulli(p)) to switch out the original values for the shuffled values. In the official custom layer guide, I see new layers inside a custom layer in either the init or the build methods of the class. As my layer needs a unique random mask at each minibatch, I have placed the mask generation in the call method of the class. The code is below:

class CutMix(tf.keras.layers.Layer):
    def __init__(self, noise):
        super(CutMix, self).__init__()
        self.noise = noise
    
    def call(self, inputs, training=None):
        if training:
            shuffled = tf.stop_gradient(tf.random.shuffle(inputs))
            msk = tf.keras.backend.random_bernoulli(inputs.shape, p=1 - self.noise, dtype=tf.float32)
            print(msk)
            return msk * inputs + (tf.ones_like(msk) - msk) * shuffled
        return inputs

Does creating this mask layer (and the shuffled layer for that matter) in the class method create a new mask layer each call, thereby exploding the computational graph's size? If this is the case, how can I incorporate randomness inside a layer (such as shuffling the minibatch or creating a mask) without this bug?

tf.keras tensorflow python
1 Answers

I need random effects at each call of my custom layer. If I create a random mask (and also create a new tensor by shuffling the input data) in the call method instead of the init or build methods, will this create new nodes on the computation graph with every function call?

Doing this, i.e., msk = tf.keras.backend.random_bernoulli(inputs.shape, p=1 - self.noise, dtype=tf.float32) in the call method will not create new nodes in the graph each time the model is called.

From their source codes:

def random_bernoulli(shape, p=0.0, dtype=None, seed=None):
  """Returns a tensor with random bernoulli distribution of values.
  Args:
      shape: A tuple of integers, the shape of tensor to create.
      p: A float, `0. <= p <= 1`, probability of bernoulli distribution.
      dtype: String, dtype of returned tensor.
      seed: Integer, random seed.
  Returns:
      A tensor.
  """
  if dtype is None:
    dtype = floatx()
  if seed is None:
    seed = np.random.randint(10e6)
  return array_ops.where_v2(
      random_ops.random_uniform(shape, dtype=dtype, seed=seed) <= p,
      array_ops.ones(shape, dtype=dtype), array_ops.zeros(shape, dtype=dtype))

showing tf.keras.backend.random_bernoulli is implemented using tf.random.uniform and tf.where. It will work properly at least in the codes you have shown, i.e., generating random(different) tensors each time the graph is executed and the graph is static. In more complex graphs, tf.random.uniform and other old RNG APIs may fail. Check https://www.tensorflow.org/guide/random_numbers for the new RNG APIs.

To be clear, do not create/init CutMix in the call method of other custom keras layers. You just call it in the call method of other custom keras layers and tf.random.uniform will generate random(different) tensors each time the graph is executed.

user_55239200
Related