How to create a custom layer for Sampling in Keras Tensorflow?

Question

I'm building a CNN in Keras with a Tensorflow backend, and I'd like to introduce a custom layer that should perform the following:

• Output a tensor of same shape and dtype as the input tensor.
• The output is made of few samples of the input tensor, let's say 25%. The rest of the output tensor should be zeros.
• The samples must be picked at random, such that the highest values pixel are sampled with higher probability . In other words, the probability distribution should be the input tensor itself (normalized).

For now, I've managed to build a mockup where I pick the top 25% pixels of the input tensor and create an output tensor of same size only from them. But it is not a random sampling.

Ideally I'd like to use a tensorflow equivalent of : np.random.choice(input_tensor, num_samples, input_tensor_normalized) where the third argument is the probability distribution to follow. Note that this only works on 1D np.array.

I've heard of tf.random.multinomial but it's depreciated and tf.random.categorical takes logits as inputs (I don't think it's my case) and doesn't propose a probability distribution.

A possibility is to reshape the input tensor as a vector, perform 1D sampling in Tensorflow if there is a way, construct a similar vector with the sampled values at the corresponding index and zeros elsewhere, and then reshape as a tensor afterwards.

Any other idea?

Should I move to PyTorch?

Thank you very Much

Answer 1

You can still use the tf.random.categorical . The logits are just the unnormalised log probabilities. So if you already have your probability distribution ready to go you can perform:

samples = tf.random.categorical(tf.log(input_tensor_normalized), num_samples)