如何用tf.keras.layers.Layer子类替代分组卷积代码中的Lambda层？

Lambda层的局限性示例 #

使用Lambda层实现缩放逻辑的代码如下：

scale = tf.Variable(1.)
scale_layer = tf.keras.layers.Lambda(lambda x: x * scale)

由于Lambda层不会直接追踪外部的scale变量，该变量不会被纳入scale_layer.trainable_weights列表，若将此层用于Keras Model中，scale将无法参与训练。

子类化Layer的优化实现 #

通过编写自定义Layer子类，可以解决上述问题，让变量被正确追踪：

class ScaleLayer(tf.keras.layers.Layer):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.scale = tf.Variable(1.)

    def call(self, inputs):
        return inputs * self.scale

分组卷积中Lambda层的替换方案 #

针对给定的分组卷积代码，其中最后一行的Lambda层用于按通道维度切片提取分组特征，我们可以编写对应的自定义Layer子类来替代：

1. 自定义通道切片Layer #

class ChannelSliceLayer(tf.keras.layers.Layer):
    def __init__(self, start_channel, end_channel, **kwargs):
        super().__init__(**kwargs)
        self.start_channel = start_channel
        self.end_channel = end_channel
    
    def call(self, inputs):
        # 假设输入为NHWC格式，通道维度在最后一位
        return inputs[:, :, :, self.start_channel:self.end_channel]

注：如果输入是NCHW格式，需要调整切片维度为inputs[:, self.start_channel:self.end_channel, :, :]

2. 修改分组卷积代码 #

将原代码中的Lambda层替换为自定义Layer：

def grouped_convolution_block(inputs, num_filters, kernel_size, strides, cardinality): 
    # 添加分组卷积块
    group_list = [] 
    grouped_channels = int(num_filters / cardinality)
    
    for c in range(cardinality):
        start = c * grouped_channels
        end = (c + 1) * grouped_channels
        x = ChannelSliceLayer(start, end)(inputs)
        # 后续可继续添加卷积等操作
        group_list.append(x)

内容的提问来源于stack exchange，提问作者joseph king