How to randomly pick and mask a portion of a Tensor in Tensorflow (python)

user_4015806 0
Viewed 808

I'm training a denoising autoencoder in Tensorflow 2, one part of the run time is spent on CPU doing masking of a portion of the input data, randomly selecting the indices to be masked, then setting their values to zero. This is my masking function, this masking is repeated on the beginning of each epoch, at different v values:

import numpy as np

def masking_noise(X, v):

    X_noise = X.copy()

    n_samples = X.shape[0]
    n_features = X.shape[1]
    v = int(np.round(n_features*v))
    for i in range(n_samples):
        mask = np.random.choice(n_features, v, replace=False)

        for m in mask:
            X_noise[i][m] = np.repeat(0.,X.shape[2])

    return X_noise

Here is a toy example:

a = np.array([[[1., 0.],
        [1., 0.],
        [1., 0.],
        [1., 0.],
        [0., 1.]],

       [[1., 0.],
        [1., 0.],
        [1., 0.],
        [1., 1.],
        [0., 1.]],

       [[1., 0.],
        [1., 0.],
        [1., 0.],
        [1., 0.],
        [1., 1.]]])

masking_noise(a, 0.40)

Output:

array([[[1., 0.],
        [0., 0.],
        [1., 0.],
        [1., 0.],
        [0., 0.]],

       [[0., 0.],
        [0., 0.],
        [1., 0.],
        [1., 1.],
        [0., 1.]],

       [[1., 0.],
        [1., 0.],
        [1., 0.],
        [0., 0.],
        [0., 0.]]])

My question is, how could I do the same masking operation in Tensorflow?

tensorflow2.0 tensorflow masking numpy python
1 Answers

I think I finally figured it out, it was easy to debug this problem with Tensorflow 2, so I was able to solve this when I changed from TF1 to TF2:

def mask_data(y_true, mask_ratio, verbose=0):

    nf = tf.cast(tf.shape(y_true)[1], tf.float32)
    mask_portion = tf.math.round( tf.math.multiply(nf,(1-mask_ratio)) )
    mask_portion = tf.cast(mask_portion, tf.int32)

    z = -tf.math.log(-tf.math.log(tf.random.uniform(tf.shape(y_true)[0:-1],0,1))) 
    _, indices = tf.nn.top_k(z, mask_portion)
    one_hots = tf.one_hot(indices, tf.shape(y_true)[1])
    mask = tf.reduce_max(one_hots, axis=1)
    mask = tf.expand_dims(mask,axis=-1)
    mask_tiles = tf.tile(mask,[1,1,tf.shape(y_true)[-1]]) 
    masked = tf.multiply(mask_tiles,toy_example)
    if(verbose>0):
        print("\nRandomly selected indices:", indices)
        print("\n2D mask (per variant)", mask)
        print("\n3D mask (per allele)", mask_tiles)
        print("\nmasked results", masked)

    return masked

Then I can run it like this:

toy_example = np.array([[[1., 0.],
    [1., 0.],
    [1., 0.],
    [1., 0.],
    [0., 1.]],

   [[1., 0.],
    [1., 0.],
    [1., 0.],
    [1., 1.],
    [0., 1.]],

   [[1., 0.],
    [1., 0.],
    [1., 0.],
    [1., 0.],
    [1., 1.]]])

mask_ratio = 0.40
result = mask_data(toy_example, mask_ratio, verbose=0)
print(result)

The result will look like this:

tf.Tensor(
[[[1. 0.]
  [1. 0.]
  [0. 0.]
  [1. 0.]
  [0. 0.]]

 [[1. 0.]
  [0. 0.]
  [0. 0.]
  [1. 1.]
  [0. 1.]]

 [[0. 0.]
  [0. 0.]
  [1. 0.]
  [1. 0.]
  [1. 1.]]], shape=(3, 5, 2), dtype=float32)
user_40158060
Related