[FIXED] Making a Sampling Layer in Keras

Issue

I am trying to build a custom Keras Layer that returns a one hot vector of a class chosen from a previous softmax layer, i.e, if Sofmax layer returns [0.4 0.1 0.5] I want to make a ramdom choice on classes 0 to 2 according to this softmax probability.

Here is what I have done so far:

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.keras import utils



def sampling(x):

    # HERE: I need to input 'x' into the 'tf.math.log' function, but after trying it created an error

    samples = tf.random.categorical(tf.math.log([[0.4 0.1 0.5]]), 1) # samples should be like [[z]] with z a number between 0 and 2

    return utils.to_categorical(samples[0][0], num_classes=3)



x = keras.Input(shape=(1,1,))

x, _, _ = layers.LSTM(100, return_sequences=True, return_state=True)(x)

x = layers.Dense(3, activation="softmax")(x)

x = layers.Lambda(sampling)(x)

This code returns:

/usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/utils/np_utils.py in to_categorical(y, num_classes, dtype) 67 68 """ ---> 69 y = np.array(y, dtype='int') 70 input_shape = y.shape 71 if input_shape and input_shape[-1] == 1 and len(input_shape) > 1:

TypeError: array() takes 1 positional argument but 2 were given

Here is the Google Colab link

Solution

You can do it with tf.one_hot.

If the input of the sampling function is 2D:

X = np.random.uniform(0,1, (100,1,1))
y = tf.keras.utils.to_categorical(np.random.randint(0,3, (100,)))

def sampling(x):
    zeros = x*0 ### useless but important to produce gradient
    samples = tf.random.categorical(tf.math.log(x), 1)
    samples = tf.squeeze(tf.one_hot(samples, depth=3), axis=1) 
    return zeros+samples


inp = Input(shape=(1,1,))
x, _, _ = LSTM(100, return_sequences=False, return_state=True)(inp)
x = Dense(3, activation="softmax")(x)
out = Lambda(sampling)(x)

model = Model(inp, out)
model.compile('adam', 'categorical_crossentropy')
model.fit(X,y, epochs=3)

If the input of the sampling function is 3D:

tf.random.categorical accepts only 2D logits. You can adapt the operation for 3D logits using tf.map_fn

X = np.random.uniform(0,1, (100,1,1))
y = tf.keras.utils.to_categorical(np.random.randint(0,3, (100,)))
y = y.reshape(-1,1,3)

def sampling(x):
    zeros = x*0 ### useless but important to produce gradient
    samples = tf.map_fn(lambda t: tf.random.categorical(tf.math.log(t), 1), x, fn_output_signature=tf.int64)
    samples = tf.squeeze(tf.one_hot(samples, depth=3), axis=1) 
    return zeros+samples


inp = Input(shape=(1,1,))
x, _, _ = LSTM(100, return_sequences=True, return_state=True)(inp)
x = Dense(3, activation="softmax")(x)
out = Lambda(sampling)(x)

model = Model(inp, out)
model.compile('adam', 'categorical_crossentropy')
model.fit(X,y, epochs=3)

here the running notebook

Answered By - Marco Cerliani