TensorFlow Conv2D层如何让所有输入通道复用同一卷积核？

默认的Conv2D层会为每个输入通道分配独立的卷积核，因此权重形状为(H, W, C, 5)。要让所有输入通道复用同一个卷积核（权重形状变为(H, W, 1, 5)），需要通过自定义层或手动卷积逻辑实现，以下是两种可行方案：

方案一：自定义共享卷积核的Conv2D层 #

通过继承tf.keras.layers.Layer，手动定义共享卷积核，并在计算时将核广播到所有输入通道：

import tensorflow as tf

class SharedConv2D(tf.keras.layers.Layer):
    def __init__(self, filters, kernel_size, activation=None, **kwargs):
        super().__init__(**kwargs)
        self.filters = filters
        self.kernel_size = kernel_size
        self.activation = tf.keras.activations.get(activation)
    
    def build(self, input_shape):
        # 定义共享卷积核：形状为(H, W, 1, filters)
        self.kernel = self.add_weight(
            shape=(*self.kernel_size, 1, self.filters),
            initializer='glorot_uniform',
            name='shared_kernel'
        )
        self.bias = self.add_weight(
            shape=(self.filters,),
            initializer='zeros',
            name='bias'
        )
        super().build(input_shape)
    
    def call(self, inputs):
        # 将共享核广播到所有输入通道，形状变为(H, W, C, filters)
        tiled_kernel = tf.tile(self.kernel, [1, 1, tf.shape(inputs)[-1], 1])
        # 执行卷积运算
        outputs = tf.nn.conv2d(inputs, tiled_kernel, strides=[1,1,1,1], padding='VALID')
        outputs = outputs + self.bias
        if self.activation is not None:
            outputs = self.activation(outputs)
        return outputs

# 测试使用
input_shape = (32, 32, 3)  # H, W, C=3
x_input = tf.keras.Input(shape=input_shape)
x_conv = SharedConv2D(5, (32, 32), activation='relu')(x_input)
model = tf.keras.Model(inputs=x_input, outputs=x_conv)

# 查看权重形状，输出为(32, 32, 1, 5)
print(model.get_weights()[0].shape)

方案二：手动实现卷积逻辑 #

直接使用tf.nn.conv2d配合核的广播操作，无需自定义层：

import tensorflow as tf

# 示例输入：(N, H, W, C) = (1, 32, 32, 3)
inputs = tf.random.normal((1, 32, 32, 3))
H, W = 32, 32
filters = 5

# 定义共享卷积核
shared_kernel = tf.Variable(tf.random.normal((H, W, 1, filters)))
bias = tf.Variable(tf.zeros((filters,)))

# 将核广播到所有输入通道
tiled_kernel = tf.tile(shared_kernel, [1, 1, tf.shape(inputs)[-1], 1])
# 执行卷积并激活
outputs = tf.nn.conv2d(inputs, tiled_kernel, strides=[1,1,1,1], padding='VALID')
outputs = tf.nn.relu(outputs + bias)

# 查看共享核形状，输出为(32, 32, 1, 5)
print(shared_kernel.shape)

说明 #

默认Conv2D的设计目标是为每个输入通道学习独立的特征模式，因此不支持直接设置通道共享卷积核。上述两种方案通过广播共享核的方式，强制所有通道使用同一组卷积参数，实现了你需要的效果。

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