TF2 Keras中如何共享卷积层权重并设置不同膨胀率共同训练？

问题根源分析 #

你遇到的RuntimeError: Cannot get value inside Tensorflow graph function是因为在模型的call方法（图模式执行）中调用了get_weights()和set_weights()——这两个方法属于Eager模式操作，只能处理numpy数组，无法在符号化的图计算中执行。同时原代码的逻辑确实存在缺陷：只有第一个卷积层的权重是可训练的，其余层的权重是手动复制的，反向传播时这些复制的权重不会产生梯度，无法实现“各膨胀率共同影响权重”的目标。

正确实现思路 #

放弃创建多个Conv2D层，改为只定义一组共享的可训练卷积权重，然后通过tf.nn.conv2d手动为每个膨胀率执行卷积操作。这样所有膨胀率的卷积计算都会复用同一组权重，反向传播时所有路径的梯度都会累积到这组权重上，自然实现各膨胀率共同影响权重的效果，同时完全兼容图模式，避免内存占用过高的问题。

修正后的代码 #

import tensorflow as tf

class customModel(tf.keras.Model):
    def __init__(self, num_filters, dmin, dmax, custom_initializer='glorot_uniform'):
        super().__init__()
        self.num_filters = num_filters
        self.kernel_size = (1, 2)
        self.custom_initializer = tf.keras.initializers.get(custom_initializer)
        # 定义共享的卷积权重和偏置（先预留输入通道维度，第一次调用时再确定）
        self.kernel = self.add_weight(
            shape=(self.kernel_size[0], self.kernel_size[1], None, num_filters),
            initializer=self.custom_initializer,
            trainable=True,
            name='shared_conv_kernel'
        )
        self.bias = self.add_weight(
            shape=(num_filters,),
            initializer='zeros',
            trainable=True,
            name='shared_conv_bias'
        )
        # 保存所有需要使用的膨胀率
        self.dilation_rates = [(1, i) for i in range(dmin, dmax+1)]

    def call(self, inputs):
        # 延迟初始化：第一次调用时确定输入通道数，补全权重的维度
        if self.kernel.shape[2] is None:
            input_channels = inputs.shape[-1]
            # 用自定义初始化器生成对应维度的权重
            init_kernel = self.custom_initializer(
                shape=(self.kernel_size[0], self.kernel_size[1], input_channels, self.num_filters)
            )
            self.kernel.assign(init_kernel)
        
        outputs = []
        # 对每个膨胀率执行卷积
        for dilation in self.dilation_rates:
            conv_out = tf.nn.conv2d(
                inputs,
                self.kernel,
                strides=(1, 1, 1, 1),
                padding='SAME',
                dilations=dilation,
                data_format='NHWC'
            )
            conv_out = tf.nn.bias_add(conv_out, self.bias)
            outputs.append(conv_out)
        
        # 拼接所有膨胀率的卷积结果
        return tf.concat(outputs, axis=-1)

关键细节说明 #

1. 共享权重定义：仅创建一组kernel和bias作为模型的可训练变量，所有膨胀率的卷积都复用这组变量。
2. 延迟初始化：由于输入通道数在模型初始化阶段无法确定，在第一次call时补全权重的输入通道维度，确保与输入数据匹配。
3. 图模式兼容：使用tf.nn.conv2d原生图操作，完全符合TensorFlow图模式的要求，不会触发Eager模式的内存问题。
4. 梯度累积：每个膨胀率对应的卷积操作产生的梯度都会自动累积到共享权重上，训练时所有膨胀率的计算都会影响权重更新，实现预期目标。

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