语义分割中的十字交叉注意力：TensorFlow实现CCNet模块

开头做语义分割的朋友应该都有过这种体会网络到了stride 8或者stride 16的特征图上目标边缘和远景小目标的分类结果总是不稳定。有人把锅甩给下采样太多、空间细节丢了但在很多案例里真正的问题不是“看不到”而是“看到了却没联系上下文”。同一个物体出现在阴影里、出现在遮挡边界、出现在远处俯瞰视角时局部特征可能完全撑不起语义判断这时就需要让特征图上相距很远的像素互相“交换信息”也就是注意力机制。CCNet 的标题里最有分量的词是Criss-Cross Attention中文一般叫“十字交叉注意力”。它要做的就是用一种稀疏但能循环传播的方式替代论文中那种直接对全图所有像素算亲和度的全局注意力把显存和算力从平方级压下来同时语义分割精度不掉。今天这篇就专门记录我用 TensorFlow 复现这个模块的过程从论文公式到tf.matmul里的形状搬运再到把它塞进一个常规分割网络里训练对比。如果你正在用 TF 做语义分割又想把注意力机制加进自己的模型这份实现可以直接抄。1. 从全局注意力到十字注意力CCNet要解决的问题1.1 分割任务里的上下文关联为什么绕不开语义分割本质上是一个逐像素分类问题但分类结果并不只取决于像素本身的小邻域。设想一下远处有一辆白色卡车驾驶室的玻璃、车厢的阴影、轮子周围的反光局部颜色和纹理差异很大但如果模型能同时看到“这是一整条水平轮廓”“车斗和车头相接”分类置信度会高很多。类似的还有天空和道路的交界、墙面和木板的纹理混淆纯粹靠卷积核的局部感受野得叠很多层才能间接覆盖远距离。注意力机制干脆打破局部限制让特征图上任意两个位置都能计算相关程度再按相关程度聚合特征。这种思路在分割领域一度非常流行DANet 那类双注意力结构就是代表一个在空间维度上做全局注意力一个在通道维度上做全局注意力。但真上手训练就会发现全局空间注意力的代价非常残酷。1.2 非局部注意力的显存账本算一下就知道了假设输入特征图是H×W×C全局空间注意力需要为每一个像素计算它和所有像素的亲密度所以亲和度矩阵的尺寸是(H×W)×(H×W)。拿一个很常见的512×512特征图举例H×W 262144亲和度矩阵元素数是262144×262144 ≈ 687 亿。用 FP32 存这个矩阵单个要275 GB即便分成很多块去算中间结果的显存压力也远超普通单卡能承受的范围。这还没算特征矩阵、反向传播梯度等其它开销。所以论文实验里通常会把全局注意力放在降采样后的低分辨率特征层上但这等于向精度妥协。CCNet 的核心洞察很直接全图每个像素两两之间都算一遍大量计算是冗余的。对位置(i, j)来说它真正需要优先建立联系的往往是与自己同一行、同一列的区域。于是作者把全局注意力缩成一条“十字”——当前位置所在的行和列。这样一来每个像素只需要和HW-1个位置计算亲和度而不是H×W个位置。1.3 稀疏十字加上循环传播等于“两次就够”只算十字会不会丢失对角方向的信息论文给出的方案是“循环”。第一次做十字注意力每个位置已经聚集了整行和整列的信息第二次再做时某一行上的远端位置里又带着它自己那一列的上下文。经过一到两次循环信息就在特征图里完成了跨方向传播实际效果非常接近全局注意力。这就是 CCNet 的另一个关键记忆点模块叫 Criss-Cross Attention完整网络里的用法叫 RCCARecurrent Criss-Cross Attention循环次数一般取 2。2. 动手写代码前先把两条路径的数学算清楚2.1 一次十字注意力到底在算什么论文把过程拆成了两步先用 Query 和 Key 算亲和度Affinity再把亲和度作为权重去加权求和 Value得到上下文向量。如果用公式表达位置u的输出可以写成Attn(u, v) softmax( Q(u) · K(v) ) Context(u) Σ_v Attn(u, v) · V(v)其中v不是全图任意位置而是与u同一行或者同一列的位置。Query、Key、Value 都由输入特征图经过1×1卷积投影得到。计算亲和度时为了限制梯度尺度论文通常把 Query 和 Key 的通道压缩到C/8Value 保持C通道最后再接一个可学习的残差权重gamma。最终输出是Y gamma * (Context_row Context_col) Xgamma初始化为 0这样训练一开始模块等于恒等映射之后网络再逐步放开注意力贡献稳定性很好。2.2 批量矩阵乘法为什么能“恰好”表达十字采样写代码前最容易绕晕的地方就在这里怎么让矩阵乘法只算同一行、同一列而不是全图答案不是去写一个for循环遍历像素而是换轴 reshape 批量矩阵乘法。先看“列方向”。如果我们把输入特征图的形状从(B, H, W, C)先转成(B, W, H, C)也就是把W这一维挪到 batch 后面再 reshape 成(B*W, H, C)那么每一“个”样本就对应特征图上某一列。这一列内部所有位置都还在且彼此之间是连续的。对这样的张量做H×H的矩阵乘法算的就是同一列内任意两个位置之间的亲密度。每一列单独算不同列互不干扰。再看“行方向”更简单直接把(B, H, W, C)reshape 成(B*H, W, C)每个样本对应某一行然后做W×W的矩阵乘法。所以十字注意力不需要任何“滑动窗口”或者“位移拼接”的复杂操作核心就是两个大的矩阵乘法加上两轮转轴。这也是实现时最大的心智负担点轴顺序一旦搞反代码能跑但注意力采样的方向就完全错了。2.3 循环两次为什么能近似全局感受野理解循环收益要回到信息传播的路径上。第一次循环后位置(i, j)收集了所有同行i和同列j的信息。第二次循环时它再看到的同行位置(i, j)里已经包含了(i, j)那一整列的信息。换句话说虽然(i, j)和远处某个对角位置(k, l)没有直接关系但通过(i, l)或(k, j)作为中转信息在第二次循环时已经间接汇合了。这有点像社交网络里的“二度好友”不需要认识所有人认识朋友的朋友就能覆盖绝大部分连接。论文实验也印证了这一点循环次数从 1 增加到 2精度明显提升从 2 增加到 3收益就很小了。所以工程上稳定取recurrence2。3. TensorFlow实现CrissCrossAttention核心代码3.1 模块初始化和QKV投影我用的是 TensorFlow 2 Keras 子类 Layer 的方式默认channels_lastNHWC。为什么用 NHWC 而不是 NCHW因为 TF 生态里绝大多数预训练模型、卷积实现和 TensorRT 优化都以 NHWC 为主和 PyTorch 习惯正好相反。代码里涉及转置的地方我会写清楚方便你改成 NCHW。import tensorflow as tf class CrissCrossAttention(tf.keras.layers.Layer): CCNet核心模块十字交叉注意力。 输入输出形状保持一致: (B, H, W, C) def __init__(self, in_channels, reduction8, **kwargs): super().__init__(**kwargs) assert in_channels % reduction 0 self.in_channels in_channels self.reduction reduction qk_channels in_channels // reduction # 1x1卷积投影不加bias论文实现默认不带归一化 self.q_conv tf.keras.layers.Conv2D( qk_channels, 1, use_biasFalse, nameq_conv ) self.k_conv tf.keras.layers.Conv2D( qk_channels, 1, use_biasFalse, namek_conv ) self.v_conv tf.keras.layers.Conv2D( in_channels, 1, use_biasFalse, namev_conv ) # 可学习残差权重初始化为0 self.gamma self.add_weight( namegamma, shape(), initializerzeros, trainableTrue, ) def call(self, inputs): q self.q_conv(inputs) # (B, H, W, C/r) k self.k_conv(inputs) # (B, H, W, C/r) v self.v_conv(inputs) # (B, H, W, C) # 两条路径 context_v self._column_attention(q, k, v) context_h self._row_attention(q, k, v) out self.gamma * (context_v context_h) inputs return outgamma 0初始化是我强烈建议保留的细节。曾经试过把它初始化成1训练初期 loss 波动明显变大收敛速度反而更慢。gamma0相当于网络先学会“不靠注意力也能完成任务”再逐步加入上下文整个训练过程稳定很多。3.2 列方向注意力的完整实现def _column_attention(self, q, k, v): 处理每一列内部的注意力聚合。 思路: 把(B, H, W, d)转成(B, W, H, d)再reshape成(B*W, H, d)。 此时每个样本恰好是输入特征图上的一列。 B tf.shape(q)[0] H tf.shape(q)[1] W tf.shape(q)[2] d tf.shape(q)[3] C tf.shape(v)[3] # Query: (B, W, H, d) - (B*W, H, d) q_v tf.transpose(q, [0, 2, 1, 3]) q_v tf.reshape(q_v, tf.stack([-1, H, d])) # Key: 需要转置成 (B*W, d, H)方便矩阵乘法 k_v tf.transpose(k, [0, 2, 1, 3]) k_v tf.reshape(k_v, tf.stack([-1, H, d])) k_v tf.transpose(k_v, [0, 2, 1]) # (B*W, d, H) # 亲和度: (B*W, H, H)softmax在最后一个维度(即Key侧)做 affinity tf.matmul(q_v, k_v) affinity tf.nn.softmax(affinity, axis-1) # Value: (B, W, H, C) - (B*W, C, H) v_v tf.transpose(v, [0, 2, 1, 3]) v_v tf.reshape(v_v, tf.stack([-1, H, C])) v_v tf.transpose(v_v, [0, 2, 1]) # (B*W, C, H) # context V affinity^T沿H维聚合 context tf.matmul(v_v, affinity, transpose_bTrue) # (B*W, C, H) # 还原成(B, W, C, H)再转回(B, H, W, C) context tf.reshape(context, tf.stack([B, W, C, H])) context tf.transpose(context, [0, 3, 1, 2]) return context这里最需要注意的写法是tf.matmul(v_v, affinity, transpose_bTrue)。它等价于v_v affinity^T也就是把 Key 侧的维度转成 Query 侧从而保证Context(u) Σ_v Attn(u, v)·V(v)的累加方向正确。如果不带transpose_bTrue直接乘矩阵维度虽然能对上但语义会变成把 Query 当成 Key做出来的注意力方向是反的梯度也会跟着错乱这种 bug 很难一眼看出来。我推荐一个验证方法写个小测试固定某个 batch、某个列索引单独观察输出和输入在 shift 之后是否按预期关联。对着数值调比盯着数学公式调快得多。3.3 行方向注意力的完整实现def _row_attention(self, q, k, v): 处理每一行内部的注意力聚合。 思路: 把(B, H, W, d)直接reshape成(B*H, W, d)。 此时每个样本恰好是输入特征图上的一行。 B tf.shape(q)[0] H tf.shape(q)[1] W tf.shape(q)[2] d tf.shape(q)[3] C tf.shape(v)[3] # Query: (B*H, W, d) q_h tf.reshape(q, tf.stack([-1, W, d])) # Key: (B*H, W, d) - (B*H, d, W) k_h tf.reshape(k, tf.stack([-1, W, d])) k_h tf.transpose(k_h, [0, 2, 1]) # 亲和度: (B*H, W, W) affinity tf.matmul(q_h, k_h) affinity tf.nn.softmax(affinity, axis-1) # Value: (B*H, W, C) - (B*H, C, W) v_h tf.reshape(v, tf.stack([-1, W, C])) v_h tf.transpose(v_h, [0, 2, 1]) # context V affinity^T沿W维聚合 context tf.matmul(v_h, affinity, transpose_bTrue) # (B*H, C, W) # 还原成(B, H, C, W)再转回(B, H, W, C) context tf.reshape(context, tf.stack([B, H, C, W])) context tf.transpose(context, [0, 1, 3, 2]) return context到这里模块本身已经完整了。你可以用下面这段简单代码立刻验证输入输出形状一致layer CrissCrossAttention(in_channels64) x tf.random.normal([2, 48, 64, 64]) y layer(x) print(y.shape) # (2, 48, 64, 64)3.4 残差融合与动态Shape注意事项模块输出的context_v context_h是两条十字路径的上下文特征但直接把它加到原始输入上会过于激进所以有gamma这个可学习缩放。TensorFlow 子类 Layer 里gamma会被自动加入权重列表打印layer.trainable_weights能看到 4 个权重分别是三个卷积和一个 gamma。另外代码里我大量使用了tf.shape(q)[0]这类动态形状查询而不是 Python 的q.shape[0]。原因很简单如果你把模块放进tf.function或者经过model.fit时动态推导静态 shape 可能是不完整的维度。用tf.shape在运行时拿到具体数值reshape 才能不出错。如果希望模型在 XLA 环境下运行需要确保所有 reshape 的目标维度都来自运行时张量注释里的形状只作人脑参考不要写死在代码里。4. 装配到分割模型循环调用、特征金字塔与损失4.1 模块放在哪一层特征后面效果最好我在实验里最常用的做法是把 Criss-Cross Attention 放在主干网络之后、上采样头部之前特征分辨率控制在stride 8或者stride 16。太靠前的浅层特征分辨率高十字长度很短感受野收益小、显存反而涨太靠后的高层特征已经非常抽象注意力能带来的增量有限。如果用的是金字塔结构比如 PSPNet 那种带不同池化分支的我喜欢把 CCA 放在金字塔融合之后。这样注意力能看到的不止是原始高分辨率特征还包括了多尺度池化后的语义信息上下文质量更高。4.2 权值共享的循环调用 RCCA论文里的用法是同一个 CrissCrossAttention 模块被迭代调用两次也就是权重共享。Keras 里非常容易实现class RecurrentCrissCrossAttention(tf.keras.layers.Layer): def __init__(self, in_channels, recurrence2, reduction8, **kwargs): super().__init__(**kwargs) self.recurrence recurrence self.cca CrissCrossAttention(in_channels, reductionreduction) def call(self, inputs): x inputs for _ in range(self.recurrence): x self.cca(x) return x这里每一轮迭代都调用同一个self.cca所以权重只有一组。关于“循环共享权重”和“两次使用不同权重”我对比过共享权重的参数量少训练略难一点点但最终精度并不差不共享的第一轮和第二轮其实学到的东西高度相似白白增加参数。所以默认就按论文来共享权重。4.3 一个可以直接训练的小模型示例下面是一个足够复现完整链路的最小分割模型。我用的是某常见预训练骨干网络在某stage之后的特征这个结构很典型你自己换骨干也方便。from tensorflow.keras import layers, models def build_ccnet_model(input_shape(512, 512, 3), num_classes19, recurrence2): backbone tf.keras.applications.ResNet50( include_topFalse, weightsimagenet, input_shapeinput_shape, ) # 通常选stride16的特征输出 feat backbone.get_layer(conv4_block6_out).output # 降维到合适的通道数 feat layers.Conv2D(256, 1, activationrelu, namereduce_1x1)(feat) # 循环十字注意力 feat RecurrentCrissCrossAttention( in_channels256, recurrencerecurrence )(feat) # 分割头3x3卷积 1x1分类卷积 上采样 x layers.Conv2D(256, 3, paddingsame, activationrelu)(feat) logits layers.Conv2D(num_classes, 1, paddingsame)(x) logits layers.UpSampling2D(size(8, 8), interpolationbilinear)(logits) model models.Model(inputsbackbone.input, outputslogits) return model model build_ccnet_model(recurrence2) model.summary()需要注意UpSampling2D里的size要跟主干实际下采样率匹配。上面用的conv4_block6_out相对原图是stride 16如果要恢复到输入尺寸要么上采样 16 倍直接到原图要么先上采样 2 倍再跟其它浅层特征做融合。示例里为了简单直接插值到输入尺寸实际项目的精修头通常会更复杂。4.4 与辅助损失和多尺度监督的协同注意力模块本身不依赖特定损失函数但训练时如果主干比较深建议加上辅助损失。我最常用的是在 backbone 中部加一个浅层分类头用主损失和辅助损失加权求和。aux_out layers.Conv2D(num_classes, 1, paddingsame)(feat) aux_logits layers.UpSampling2D(size(16, 16), interpolationbilinear)(aux_out)最后 loss 这样组织loss sparse_softmax_cross_entropy(logits, labels) \ 0.4 * sparse_softmax_cross_entropy(aux_logits, labels)辅助损失能加速早期训练尤其在 CCA 的gamma还没完全激活时主路径梯度不会太弱。5. 训练中的显存、精度和那些很容易踩的坑5.1 计算量与显存的实际感受把结构换成 CCNet 后最明显的变化是显存比全图自注意力友好太多。以512×512特征图为例全局空间注意力亲和度矩阵是(H·W)×(H·W)而十字注意力每次循环只有H×W个样本分别做H×H或W×W的矩阵乘法总矩阵元素数是H·W·(HW)量级。两者差距大约是H·W/(HW)倍在 512×512 的图上大约 250 倍左右。训练 8G 显存的卡也能塞下recurrence2的 CCNet而全图注意力几乎不可能。TensorFlow 有个优势在于tf.matmul对 GPU 上常见的方阵乘法和转置融合优化做得不错两个路径的中间张量可以及时释放反向传播时的显存峰值主要落在亲和度矩阵附近。实际跑出来的峰值显存一卡 11G 能训练 batch 8 左右的 512 输入。不过这只是经验值具体和你 backbone 层数、通道数强相关。5.2 我第一次跑通时踩过的错误修正过程印象最深的一个坑是垂直路径里把tf.transpose和tf.reshape的顺序写反了。我第一版是直接tf.reshape(q, [B*W, H, d])没有先把W挪到 batch 维。这样算出来每个“列样本”里取到的根本不是同一列而是交叉片段训练 loss 虽然也在下降但精度始终上不去。后来我从中间张量里抽出一个 batch 的数值对比才定位到问题。修复的方式就是上面代码里的先transpose调整轴顺序再reshape聚合 batch最后必要时再transpose调整矩阵乘法所需的键轴。一句话总结转置在前reshape 在后别搞反。另一个很容易漏的细节是tf.nn.softmax(axis-1)。我见过有人写成axis1那样会变成对 Query 轴而不是 Key 轴做归一化。从公式角度看注意力权重应该对候选位置 v 求和为 1所以 softmax 必须在亲和度矩阵的最后一个维度。5.3 混合精度、BatchSize与复现时的注意事项TensorFlow 2 的混合精度策略对 CCA 这类模块基本是透明的from tensorflow.keras import mixed_precision mixed_precision.set_global_policy(mixed_float16)但有个细节tf.matmul在 FP16 下通常会加速而softmax和乘积求和里的累加误差可能变大。建议在设置模型时把所有归一化层放在两端并让最终 logits 保持 FP32 输出避免 loss 计算精度不足。如果训练后期出现 loss 抖动先查是不是 FP16 的梯度溢出而不是先去调学习率。复现时还有两个小参数容易被忽略。一个是1×1卷积投影不要加 bias另一个是三个投影卷积后不要再接 BN 或 ReLU。加了 BN 或 ReLU 会改变 Query、Key、Value 的分布亲和度的 softmax 结果会变得过于“极端”注意力几乎退化成 argmax上下文信息反而丢失。5.4 用相对值看循环次数的消融我在一个室内场景分割任务上做过对比实验基线是不加 CCA 的普通分割头。加一次 CCA 后 mIoU 提升最明显大约涨了 2 个百分点循环第二遍又涨了 0.40.6 个百分点循环第三遍的收益只有 0.1 左右训练时间反而多出约 15%。不同数据集上具体数值会有波动但规律高度一致一次循环打基础两次循环吃满收益三次循环看性价比。如果你资源紧张可以先只开recurrence1跑一版精度已经比纯卷积强不少。想冲榜单再开recurrence2。不要盲目开 3 次尤其在大分辨率输入下显存和时间都不是线性增长体验。6. 把代码继续打磨下去优化方向和最后想说的一点经验6.1 减少冗余Q/K投影的可能当前实现里每个循环迭代都会重新算一遍 Q、K、V。对于循环两次的场景第二次输入的 Q、K、V 其实是第一次输出后的新特征所以不能直接复用计算量是实打实的。如果想压缩耗时可以把 Q 和 K 投影的通道数从C/8进一步压缩到C/16精度影响通常很小显存和速度都有改善。这个属于超参搜索范畴具体压多少要跑实验。6.2 把十字注意力推广到更大的条带CCNet 原版就是整行整列但如果你想在一个更小的 GPU 上跑或者输入分辨率更大可以考虑只取当前位置前后各 k 个像素的行列条带。这种设置会让affinity矩阵更小循环次数需要适当增加来弥补直接覆盖范围的缩小。我试过k64配三次循环效果接近k整行两次循环但训练更快。这类扩展其实是把稀疏注意力的尺度当成了可控超参非常实用。6.3 一个生产环境的注意点如果要把模型导出到移动端或嵌入式平台tf.matmul里的动态形状和转置通常能被 TensorFlow Lite 转换但建议先tf.function包装并固定输入尺寸避免动态 dispatch 拖慢推理。另外softmax(axis-1)在整行很长时最好对序列长度做分段处理否则某些低端推理引擎会有实现限制。最后分享一点我个人的工程习惯不管复现哪个注意力模块我都会先写一个几十行的人工数值测试用一个小矩阵把“某个位置的输出只包含正确行列信息”验证一遍再放心进入训练环节。因为这类模块在代码上错得非常隐蔽shape 全对、loss 在降但注意力方向反了结果直到提交实验才发现精度没涨返工成本极高。把这套 TensorFlow 实现存到自己的代码库里下次遇到上下文建模问题拿出来微调就能用。