创建了一个新的层来添加通道：

class ChannelAdd(tf.keras.layers.Layer):
    def __init__(self,offset=3, **kwargs):
        super(ChannelAdd, self).__init__(**kwargs)
        self.offset = offset
    def call(self, inputs):
        s_ch = inputs + tf.roll(inputs, shift=-self.offset, axis=-1) 
        x = s_ch[:,:,:,::6]
        y = s_ch[:,:,:,1::6]
        z = s_ch[:,:,:,2::6]
        concat = tf.concat([x[...,tf.newaxis], y[...,tf.newaxis], z[...,tf.newaxis]], axis=-1)
        concat = tf.reshape(concat, [tf.shape(inputs)[0],tf.shape(inputs)[1],tf.shape(inputs)[2],-1])
        return concat

型号：

layer3 = keras.Input(shape=(8,8,768))

xyz = ChannelAdd()(layer3)

model = keras.Model(inputs=layer3, outputs=xyz)

测试模型：
创建一个易于测试的特殊输入数组，其中通道包含通道号：通道0包含全0，通道1包含全1，依此类推

a = tf.zeros(shape=(5,8,8,1))
inputs = tf.concat([a,a+1,a+2], axis=-1)
for i in range(255):
    inputs = tf.concat([inputs, a+3*(i+1),a+3*(i+1)+1,a+3*(i+1)+2], axis=-1)

输出量：

outputs = model(inputs)
outputs.shape
#[5, 8, 8, 384]

outputs[:,:,:,:6]
#check first 6 channels
array([[[[ 3.,  5.,  7., 15., 17., 19.],
     [ 3.,  5.,  7., 15., 17., 19.],
     [ 3.,  5.,  7., 15., 17., 19.],
     ...,

现在，您所做的是，就我所理解的问题而言，您正在添加相同的通道，如（G+G+G+G+...）然后（R+R +R+R+...）等等...我做的代码和你做的几乎一样。但是，使用不同的策略，首先，我们需要根据3个通道拆分输入，就像我们看到的那样（R，G，B），然后我们将相应地添加通道（R-〉R，G-〉G，B-〉B）...我已经编写了一个代码，它的功能与您的代码几乎相同，但没有for循环...
第一章，绪论。

一体化GO用于在Keras.layer中实施

#for your model in one go 

split_layer = tf.keras.layers.Lambda(lambda x : tf.split(x, 256, axis=-1))(layer3)
add_layer = tf.keras.layers.Add()([tf.cast(split_layer[::2], 
                                           dtype=tf.float32), 
                                   tf.cast(split_layer[1::2], dtype=tf.float32)])
add_layer = tf.keras.layers.Lambda(lambda x : tf.split(x, 3 , axis=-1))(add_layer)
layer3 = tf.keras.layers.concatenate(add_layer, axis=0)
layer3 = tf.keras.layers.Lambda(lambda x : tf.transpose(x , perm=[1,2,3,0,4])[:,:,:,:,0])(layer3)

输出量：

<KerasTensor: shape=(None, 80, 80, 384) dtype=float32 (created by layer 'lambda_1')>

For 2 Iterations您的方法花费了23 secs，而我的方法花费了**1.7**秒。