TensorFlow子类化模型创建中层的复用性疑问

问题背景 #

在学习TensorFlow子类化创建模型时，发现教程中配置相同的BatchNormalization()和MaxPool2D()被重复定义，代码示例如下：

#Defining the class
class FeatureExtractor(Layer):
    def __init__(self):
        super().__init__()

        self.conv_1 = Conv2D(filters = 6, kernel_size = 4, padding = "valid", activation = "relu")
        self.batchnorm_1 = BatchNormalization()
        self.maxpool_1 = MaxPool2D(pool_size = 2, strides=2)

        self.conv_2 = Conv2D(filters = 16, kernel_size = 4, padding = "valid", activation = "relu")
        self.batchnorm_2 = BatchNormalization()
        self.maxpool_2 = MaxPool2D(pool_size = 2, strides=2)


    def call(self, x):
        x = self.conv_1(x)
        x = self.batchnorm_1(x)
        x = self.maxpool_1(x)

        x = self.conv_2(x)
        x = self.batchnorm_2(x)
        x = self.maxpool_2(x)

        return x

#Calling and using the class
feature_extractor = FeatureExtractor()

func_input = Input(shape=(IMG_SIZE, IMG_SIZE, 3), name="Input_Image")

x = feature_extractor(func_input)

查阅资料后得知层调用时会适配输入维度并固定，但有回复指出MaxPool2D实例可复用，BatchNormalization不可，想明确：是否部分层能适配输入维度、部分不能？是否有属性可在文档中查询该行为？

核心原因：层的状态属性差异 #

TensorFlow的层可以分为两类，这直接决定了它们是否能被复用：

1. 无状态层（可复用，自动适配输入维度） #

这类层的计算仅依赖输入数据本身，不维护任何可训练参数或运行时统计量，比如MaxPool2D、AveragePool2D、Activation、Dropout（非训练阶段）等。

• 调用时会自动适配输入的空间维度（高度、宽度）、通道数，只要操作逻辑匹配（比如池化窗口不超过输入尺寸），同一个实例可以反复调用处理不同维度的输入。
• 比如MaxPool2D的pool_size是固定的，但每次调用时会根据当前输入的H/W动态计算池化窗口的位置，不需要固定输入维度，也没有需要保存的状态信息。

2. 有状态层（不可复用，维度固定） #

这类层会维护可训练参数（如缩放系数gamma、偏移系数beta）或运行时统计量（如BatchNorm在训练阶段累计的均值、方差），比如BatchNormalization、LayerNormalization、Dense、Conv2D等。

• 第一次被调用时，层会根据输入的特征维度（比如BatchNorm默认对应输入的最后一维通道数）初始化参数和统计量的维度，之后这些维度就被固定。
• 比如例子中conv_1输出通道是6，batchnorm_1会初始化对应6通道的gamma、beta和统计量；conv_2输出通道是16，必须用新的batchnorm_2实例，复用batchnorm_1会因参数维度不匹配报错。

如何在文档中查询该属性 #

每个TensorFlow层的官方文档都会明确标注其是否为有状态（Stateful），也可以通过以下方式判断：

• 查看层的trainable_weights属性：无状态层的该属性通常为空列表，有状态层则包含可训练参数。
• 检查是否存在moving_mean、moving_variance这类运行时统计属性：比如BatchNormalization就有这些属性，而MaxPool2D没有。

复用无状态层的示例 #

将代码中的MaxPool2D实例复用，不会产生任何问题：

class FeatureExtractor(Layer):
    def __init__(self):
        super().__init__()
        self.conv_1 = Conv2D(filters=6, kernel_size=4, padding="valid", activation="relu")
        self.batchnorm_1 = BatchNormalization()
        self.maxpool = MaxPool2D(pool_size=2, strides=2)  # 仅定义一次

        self.conv_2 = Conv2D(filters=16, kernel_size=4, padding="valid", activation="relu")
        self.batchnorm_2 = BatchNormalization()

    def call(self, x):
        x = self.conv_1(x)
        x = self.batchnorm_1(x)
        x = self.maxpool(x)  # 第一次调用，适配conv_1的输出维度

        x = self.conv_2(x)
        x = self.batchnorm_2(x)
        x = self.maxpool(x)  # 复用实例，自动适配conv_2的输出维度
        return x

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