如何在TF/Keras中对张量的每个切片执行独立2D卷积操作

核心思路：用分组卷积实现无循环的独立切片卷积 #

你需要的功能可以通过Keras内置的分组卷积特性实现，不需要循环遍历，计算效率高且权重完全可训练，具体逻辑如下：

• 首先对输入张量做维度重排：将切片维度移动到最后一维（通道维度位置）
• 调用Conv2D时设置groups参数等于切片总数，此时每个切片会被分配独立的卷积权重，互相不共享
• 卷积完成后再把维度调整回你需要的格式即可

你之前尝试的两种方案的问题：

• map_fn方案需要为每个切片单独初始化卷积层，不仅速度慢，还会因为动态生成层导致权重无法正确绑定到自定义层的可训练参数列表
• DepthwiseConv2D本质是groups=输入通道数且filters=输入通道数的特殊分组卷积，只能实现每个输入通道输出1个滤波器的场景，无法支持每个切片输出多滤波器的需求，因此第二次调用不符合预期

自定义层实现代码 #

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers

class PerSliceConv2D(layers.Layer):
    def __init__(self, filters_per_slice=1, kernel_size=3, activation=None, padding="same", **kwargs):
        super().__init__(**kwargs)
        self.filters_per_slice = filters_per_slice
        self.kernel_size = kernel_size
        self.activation = keras.activations.get(activation)
        self.padding = padding
    
    def build(self, input_shape):
        # input_shape 格式为 (batch_size, num_slices, height, width)
        self.num_slices = input_shape[1]
        # 总输出通道数 = 切片数 * 每个切片输出的滤波器数
        total_filters = self.num_slices * self.filters_per_slice
        self.conv = layers.Conv2D(
            filters=total_filters,
            kernel_size=self.kernel_size,
            padding=self.padding,
            groups=self.num_slices, # 核心参数：分组数等于切片数，每个组独立卷积
            activation=self.activation
        )
        super().build(input_shape)
    
    def call(self, inputs):
        # 输入维度: (B, S, H, W) B是批量，S是切片数，H/W是高宽
        # 维度转置为 (B, H, W, S) 适配Conv2D的通道在后格式
        x = tf.transpose(inputs, perm=(0, 2, 3, 1))
        # 分组卷积，每个S对应的通道组独立计算
        x = self.conv(x)
        # 输出维度是 (B, H, W, S*K) K是每个切片输出的滤波器数
        # 转置回 (B, S, K, H, W)
        x = tf.transpose(x, perm=(0, 3, 1, 2))
        # 如果每个切片只输出1个滤波器，可挤压掉多余维度
        if self.filters_per_slice == 1:
            x = tf.squeeze(x, axis=2)
        return x

验证测试 #

你可以用以下代码验证功能是否符合预期：

# 测试输入形状：批量2，500个切片，每个100x100
test_input = tf.random.normal((2, 500, 100, 100))
layer = PerSliceConv2D(filters_per_slice=1, kernel_size=3)
output = layer(test_input)
print(output.shape) # 输出应为 (2, 500, 100, 100)，和输入形状一致
# 查看可训练参数数量：每个切片对应3x3的卷积核+偏置，总参数=500*(9+1)=5000，符合预期
print("可训练参数数量：", layer.count_params())

方案优势 #

• 完全符合需求：每个切片对应独立的可训练卷积权重，互相不共享
• 没有循环逻辑，调用底层优化的卷积算子，运行效率远高于map_fn遍历的方案
• 原生支持Keras的训练、保存、部署流程，没有兼容性问题
• 支持每个切片输出多个滤波器，只需要修改filters_per_slice参数即可，适配不同场景需求

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