创建自定义图层/函数以重新排列图层值的正确方法（带有Tensorflow的Keras）

Question

I need to rearrange a tensor values and then reshape it in Keras, however I am struggling with the proper way to to rearrange a tensor in Keras with Tensorflow backend.

This custom layer/function will iterate through the values, and then rearrange the values via a stride formula This doesn't seem to have weights, so I am assuming stateless and won't affect back propagation.

It requires list slicing though:

out_array[b,channel_out, row_out, column_out] = in_array[b,i,j,k] and this is just one of the components I am struggling with.

Here is the function/layer

def reorg(tensor, stride):

    batch,channel, height, width = (tensor.get_shape())
    out_channel = channel * (stride * stride)
    out_len = length//stride
    out_width = width//stride

    #create new empty tensor  
    out_array = K.zeros((batch, out_channel, out_len, out_width))

    for b in batch:    
        for i in range(channel):
            for j in range(height):
                for k in range(width):
                    channel_out = i + (j % stride) * (channel * stride) + (k % stride) * channel
                    row_out = j//stride
                    column_out = k//stride
                    out_array[b,channel_out, row_out, column_out] = K.slice(in_array,[b,i,j,k], size = (1,1,1,1))


    return out_array.astype("int")

I don't have much experience creating custom functions/layers in Keras, so not quite sure If I am on the right track.

Here is what the code bit is doing depending on the stride (here it's 2):

https://towardsdatascience.com/training-object-detection-yolov2-from-scratch-using-cyclic-learning-rates-b3364f7e4755

Answer 1

When you say re-arrange, do you mean change the order of your axes? There is a function called tf.transpose which you can use inside a custom layer. There is also tf.keras.layers.Permute which can be used without any custom code to re-order a tensor.

If you are asking how you can create a custom layer, there are some methods you'll need to implement. The docs explain it pretty well here: Custom Layers

from tensorflow.keras import layers
import tensorflow as tf

class Linear(layers.Layer):

  def __init__(self, units=32):
    super(Linear, self).__init__()
    self.units = units

  def build(self, input_shape):
    self.w = self.add_weight(shape=(input_shape[-1], self.units),
                             initializer='random_normal',
                             trainable=True)
    self.b = self.add_weight(shape=(self.units,),
                             initializer='random_normal',
                             trainable=True)

  def call(self, inputs):
    return tf.matmul(inputs, self.w) + self.b