Deleting tensor slices as part of a TensorFlow layer

user_13810097 0
Viewed 209

Versions: tensorflow 2.3.0, numpy 1.18.5, python 3.8.2

I'd like to remove some selected slices of an input tensor using the first layer in my TensorFlow model. For example, I have an input shape of (180, 90, 25) (where the 180 is the batch size) and I'd like to remove the list of indices indices = [3, 4, 5, 6, 7, 22, 23, 24] from the last dimension so that, after calling this layer on the input tensor, I'd get a tensor of shape (180, 90, 25 - len(indices)), where each of the selected (180, 90)-shaped tensor slices have been removed by indexing the last dimension.

Currently, I'm using this layer:

class RemoveSelectedIndices(tf.keras.layers.Layer):
    def __init__(self, indices=[3,4,5,6,7,22,23,24]):
        super(RemoveSelectedIndices, self).__init__(name="RemoveSelectedIndices")
        self.indices = self.add_weight(name="indices", shape=len(indices), dtype=tf.int32, trainable=False,
                                       initializer=lambda *args, **kwargs: indices)

    def build(self, input_shape):
        pass

    def call(self, input_tensor):
        X = tf.unstack(input_tensor, num=input_tensor.shape[-1], axis=2) # list of 25 (180, 90)-shaped slices
        indices = sorted(list(self.indices.value().numpy()))
        for i in reversed(indices):
            del X[i]
        X = tf.stack(X, axis=2) # restacking the list back together
        return X

This works perfectly fine when I'm testing it (by creating a numpy array and using tf.convert_to_tensor and then calling the layer on the tensor), but when I try to build a model using this layer as the first layer, I get an error:

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

inputs = Input(shape=(90, 25))
X = RemoveSelectedIndices()(inputs)
# gives me AttributeError: 'Tensor' object has no attribute 'numpy'
# references the line indices = sorted(list(self.indices.value().numpy()))

Why is this happening, and is there any way I can get around it?

(Note: I know I could just do this to the data itself, but the dataset is huge, and I'd rather not mess with the dataset too much unless I have to.)

Thank you in advance!

keras tensorflow deep-learning machine-learning python
1 Answers

for this kind of operation you can simply use a Lambda layer

to_remove = [3,4,5,6,7,22,23,24]

def select(X, to_remove):
    X = tf.stack([X[...,i] for i in range(X.shape[-1]) if i not in to_remove], -1)
    return X

inputs = Input(shape=(90, 25))
x = Lambda(lambda x: select(x, to_remove))(inputs)
user_103750490
Related