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

我需要重新排列一个张量值，然后在Keras中重新排列它，但是我正在努力找到正确的方法来使用Tensorflow后端重新排列Keras中的张量

此自定义层/函数将遍历这些值，然后通过步长公式重新排列这些值 这似乎没有权重，所以我假设它是无状态的，不会影响反向传播

但它需要列表切片：

out\u数组[b，channel\u out，row\u out，column\u out]=in\u数组[b，i，j，k]

这只是我正在努力解决的问题之一

这是函数/层

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")

我没有太多在Keras中创建自定义函数/层的经验， 所以我不太确定我是否在正确的轨道上

下面是代码位根据步幅所做的操作（这里是2）：

---

当你说“重新排列”时，你的意思是改变轴的顺序吗？有一个名为的函数，可以在自定义层中使用。还有一个可以在没有任何自定义代码的情况下用于对张量重新排序的函数

如果您询问如何创建自定义层，则需要实现一些方法。文档在这里解释得很好：

---

请你修改一下你的问题，解释一下逻辑或者这个“重新安排”好吗？也许可以使用向量/矩阵运算来完成。@今天很抱歉，我添加了一张图片。我添加了一张图片来显示重新排列。是否需要创建完整的图层？我的印象是，只有当我需要节省重量时，层才是必要的，所以我希望我可以做一些定制的重新排列，在Tensorflow中进行切片。

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