Is there better way to repeat layers inside a tensor than using Reshape?

Viewed 39

I want to make a tensor of shape [Batch Size, 1] into [Batch Size, Channel Size]. To do so, I have written the following algorithm:

@tf.function
def fit_channel_layers(input_tensor, target_tensor):
    final_channel_layers = tf.shape(target_tensor)[-1]
    input_tensor = tf.repeat(input_tensor, repeats=[final_channel_layers])
    input_tensor = tf.reshape(input_tensor, [local_batch_size, final_channel_layers])
    return input_tensor

input_tensor = tf.convert_to_tensor([[1.],[0.],[0.],[1.],[0.],[0.]])
target_tensor = tf.convert_to_tensor([[1.,1.,1.,1.,1.,1.],
 [0.,0.,0.,0.,0.,0.],
 [0.,0.,0.,0.,0.,0.],
 [1.,1.,1.,1.,1.,1.],
 [0.,0.,0.,0.,0.,0.],
 [0.,0.,0.,0.,0.,0.]])
fit_channel_layers(input_tensor, target_tensor)

The output is appropriate

array([[1., 1., 1., 1., 1., 1.],
       [0., 0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0., 0.],
       [1., 1., 1., 1., 1., 1.],
       [0., 0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0., 0.]], dtype=float32)>

However, I'm uncertain about the performance implications of using reshape. I know with numpy it is possible to easily expand a numpy array with using comma and a multiplier operator. However, I'm wondering if there's a better way than currently written in @tf.function graph code to get this task accomplished?

0 Answers
Related