Creating a custom keras layer that can randomly select k length from the first dimension of the input tensor

Question

I'm trying to create a custom layer in Keras that each time picks k random samples from the inputs tensor.

Something like that:

import random

class RandomKAggregator(tf.keras.layers.Layer):
    def __init__(self, **kwargs):
        super(RandomKAggregator, self).__init__(**kwargs)
        self.k = 3

    def call(self, inputs):
        return inputs[[random.randint(0, len(inputs) - 1) for _ in range(self.k)]]

The goal is to pick k samples across the first dimension. For example, if the inputs tensor's shape is [500, 32, 32, 3] it should return a tensor of shape [k, 32, 32, 3], where each call for the layer it should pick different k elements.

The above implementation returns TypeError: list indices must be integers or slices, not list. I tried to use tf.gather and tf.gather_nd but couldn't resolve that.

• What is the right way to achieve the desired result?
• Also, would like to understand how TensorFlow handles indexing and how should I use lists of indices with tensors, say I want to pick [6,3,1,0] from my_tensor, and why my_tensor[[6,3,1,0]] doesn't work for me?

Answer 1

You need to know about slicing (link_1, link_2):

>>> tns = tf.random.uniform((3,2,2,3))
>>> tns
<tf.Tensor: shape=(3, 2, 2, 2), dtype=float32, numpy=
array([[[[0.2313137 , 0.17192566],
         [0.25301564, 0.20687258]],
        [[0.5184531 , 0.49340045],
         [0.41016984, 0.6423464 ]]],
       [[[0.57849526, 0.8964175 ],
         [0.86068404, 0.28210032]],
        [[0.96660316, 0.66522324],
         [0.74370325, 0.2124871 ]]],
       [[[0.24575269, 0.7576513 ],
         [0.6213573 , 0.80739546]],
        [[0.79363906, 0.16595817],
         [0.42819571, 0.05265415]]]], dtype=float32)>

>>> tns[1:3, ...]
# Or
>>> tns[1:3, :, :, :]
<tf.Tensor: shape=(2, 2, 2, 2), dtype=float32, numpy=
array([[[[0.57849526, 0.8964175 ],
         [0.86068404, 0.28210032]],
        [[0.96660316, 0.66522324],
         [0.74370325, 0.2124871 ]]],
       [[[0.24575269, 0.7576513 ],
         [0.6213573 , 0.80739546]],
        [[0.79363906, 0.16595817],
         [0.42819571, 0.05265415]]]], dtype=float32)>

Example Code: (for generating random number you can use numpy.random.randint(start, end, num_number)

import tensorflow as tf
import numpy as np

class RandomKAggregator(tf.keras.layers.Layer):
  def __init__(self, k=3):
    super(RandomKAggregator, self).__init__()
    self.k = k


  def call(self, inputs):
    rnd = np.random.randint(0, len(inputs), 1)[0]
    return inputs[rnd:rnd+self.k, ...]

layer = RandomKAggregator(k=10)
layer(tf.random.uniform((500,32,32,3))).shape

Output:

TensorShape([10, 32, 32, 3])

Edit, How pick k random elements: (Ask in comments)

• We can use tf.gather_nd and return k random selected.

import tensorflow as tf
import numpy as np

class RandomKAggregator(tf.keras.layers.Layer):
  def __init__(self, k=3):
    super(RandomKAggregator, self).__init__()
    self.k = k


  def call(self, inputs):
    rnd = np.random.randint(0, len(inputs), self.k)
    return tf.gather_nd(inputs, rnd.reshape(-1,1))

layer = RandomKAggregator(k=10)
layer(tf.random.uniform((500,32,32,3))).shape
# TensorShape([10, 32, 32, 3])