TensorFlow子类化模型中__init__内的层实例能否复用？

原始可运行代码 #

#Defining the class
class FeatureExtractor(Model):
    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)

修改后的报错代码 #

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

        self.conv_1 = Conv2D(filters = 6, kernel_size = 4, padding = "valid", activation = "relu")
        #Defining batchnorm and maxpool only once
        self.batchnorm = BatchNormalization()
        self.maxpool = MaxPool2D(pool_size = 2, strides=2)

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


    def call(self, x):
        x = self.conv_1(x)
        x = self.batchnorm(x)
        x = self.maxpool(x)

        x = self.conv_2(x)
        x = self.batchnorm(x)
        x = self.maxpool(x)

        return x

用户疑问 #

修改后运行出现维度错误，原本以为__init__中的层实例可以复用，是否因为层被调用后会适配输入维度并固定，导致后续调用出错？

解答 #

你的猜测完全正确，问题核心在于TensorFlow的层实例是有状态的，首次调用后会固定输入相关的参数维度，无法直接复用在不同维度的输入上。

具体原因：

• BatchNormalization层：首次调用时会根据输入的通道数（最后一维）初始化均值、方差等可训练参数，这些参数的维度是固定的。第一次调用它处理的是conv_1输出的6通道特征图，参数维度就固定为6；第二次再用它处理conv_2输出的16通道特征图时，维度不匹配，直接触发错误。
• MaxPool2D层：虽然它没有可训练参数，但TensorFlow的层实例在首次调用后会记录输入的空间维度信息，后续调用如果输入的空间维度变化，也可能引发维度不兼容的问题。

需要明确：TensorFlow中每个层实例只能对应模型中的一个固定位置，哪怕结构完全相同的层，也需要单独创建实例。如果想减少重复代码，可以在__init__里封装一个生成重复层组合的方法，示例如下：

class FeatureExtractor(Model):
    def __init__(self):
        super().__init__()
        self.conv_block_1 = self._make_conv_block(6)
        self.conv_block_2 = self._make_conv_block(16)

    def _make_conv_block(self, filters):
        return tf.keras.Sequential([
            Conv2D(filters=filters, kernel_size=4, padding="valid", activation="relu"),
            BatchNormalization(),
            MaxPool2D(pool_size=2, strides=2)
        ])

    def call(self, x):
        x = self.conv_block_1(x)
        x = self.conv_block_2(x)
        return x

这种写法既避免了代码重复，又保证每个卷积块里的层都是独立实例，不会出现维度冲突。

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