TensorFlow中NHWC张量reshape为[N*C,H,W,1]实现逐通道卷积问题

1. 交错问题成因 #

你的分析完全正确。NHWC格式的张量采用通道优先连续存储规则，内存排布顺序为N -> H -> W -> C，也就是每个空间位置(h,w)的C个通道值在内存中是紧挨着的，之后才是下一个空间位置、下一行、下一个样本。
reshape操作仅改变张量的维度视图，不会调整底层内存的存储顺序。你直接将[N,H,W,C]reshape为[N*C,H,W,1]时，原来属于同一个空间位置不同通道的值，会被当成新的“伪batch”中不同空间位置的单通道值，自然就出现了通道混入空间维度的交错现象。

2. 无transpose的reshape方案不存在 #

reshape的核心特性是不修改底层数据的存储顺序，而你的需求本质是要把原本在最后一维的通道维度，移动到和batch维度相邻的位置，必须通过重排内存数据的transpose操作才能实现，不存在仅靠reshape就能达成目标的方法。

3. 更高效的同核逐通道卷积实现方案 #

给你两个无需reshape、transpose，全程兼容NHWC格式的方案，性能远高于你现在的transpose版本：

方案1：使用分组卷积（最推荐） #

TensorFlow的Conv2D支持groups参数，将其设为输入通道数C时，输入会被拆分为C个独立的通道组，每个组复用同一个卷积核计算，完全符合你“同一个2D滤波器作用于所有输入通道”的需求：

_, self.h, self.w, self.c = inputs.shape
# groups等于输入通道数，每个通道组复用同一组卷积核，每组输出2个结果
self.conv = tf.keras.layers.Conv2D(filters=2*self.c, kernel_size=3, strides=1, padding='same', groups=self.c)
x = self.conv(inputs) # 输出形状为[N, H, W, 2*C]
# 拆分输出通道维度，取每个通道对应两个结果的argmax
x = tf.reshape(x, [-1, self.h, self.w, self.c, 2])
output = tf.math.argmax(x, axis=-1) # 最终输出形状[N, H, W, C]，符合要求

该方案没有任何冗余的维度转换操作，性能最高。

方案2：权重共享的深度卷积 #

你也可以用DepthwiseConv2D实现，手动强制所有通道的卷积核权重一致即可，同样不需要维度转换：

_, self.h, self.w, self.c = inputs.shape
self.conv = tf.keras.layers.DepthwiseConv2D(kernel_size=3, strides=1, padding='same', depth_multiplier=2)
# 初始化后将所有通道的卷积核替换为同一个共享核
kernel, bias = self.conv.get_weights() # kernel原始形状为[3,3,C,2]
shared_kernel = tf.tile(kernel[:,:,0:1,:], [1,1,self.c,1]) # 取第一个通道的核复制C份
self.conv.set_weights([shared_kernel, bias])
x = self.conv(inputs) # 输出形状为[N,H,W,2*C]
x = tf.reshape(x, [-1, self.h, self.w, self.c, 2])
output = tf.math.argmax(x, axis=-1)

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