TransUnet 与 CBAM 结合：优化车道线语义分割的 Transformer 注意力增强方案

简介面向高速道路车道线语义分割的TransUnet改进实现在Transformer块中集成CBAM模块通过通道注意力与空间注意力联合建模有效增强网络对细长车道线的特征表达提升分割精度。资源附带完整的高速道路车道线数据集已划分训练集与验证集代码中包含数据加载与训练流程修改train脚本参数即可直接运行默认训练100轮学习率采用余弦退火从0.01衰减到0.00001优化器为AdamW并支持调整img-size以适配不同输入尺度。评估部分输出dice、iou、recall、precision、f1、pixel accuracy等指标训练与验证结果自动保存至runs目录下的json文件推理阶段会遍历inference/img下所有图像分别生成gt阈值图像和原图叠加掩膜效果图便于逐样本检查预测质量。整个资源包共2000个文件其中1490个png和492个jpg为车道线图像及对应标注16个py为训练、验证、推理和评估脚本另有txt与readme说明文件压缩包约196.24MB目前已有148人学习下载。读者可在此基础上直接开展对比实验若需更换其他注意力机制只需将CBAM实现替换即可也可借助清晰的数据划分和评估流程快速完成毕业设计或论文实验。1. 为什么要在 TransUnet 的 Transformer 块里加 CBAM在 TransUnet 的 Transformer 块里插 CBAM网上最常见的做法是当即插即用模块往注意力层后面一放就跑数据。但对车道线这种细长目标真正决定改进是否有效的是两个问题CBAM 的通道与空间两个分支分别作用在哪个维度上以及插入后 token 序列和图像特征图之间怎样对齐。TransUnet 让 Transformer 块负责 patch 间的长距离关系CBAM 的空间注意力又擅长把响应集中到局部关键区域两者天然互补。车道线在前景里往往只占不到 2% 的像素也正是 CBAM 能补短板的场景。下面会给出不依赖外部库的 CBAM 实现、可直接替换的 Transformer 块以及数据转 mask 到评估 F1 的完整路径。2. 动手前先定实现方案CBAM 模块怎么装进 Transformer Block2.1 先分清 CBAM 两条分支分别作用在哪个维度上CBAM 由两个串联的子模块构成通道注意力ChannelGate和空间注意力SpatialGate。通道注意力对全局平均池化和全局最大池化后的结果做共享 MLP 映射得到每个通道的权重空间注意力则在通道维上取平均值与最大值拼成两通道特征后用一个 7×7 卷积输出空间权重。整个过程是一次轻量的重新校准乘法作用在特征图的不同维度上不会像 FFN 那样改变整个表达空间。原版 TransUnet 的中途计算基本都沿用 ViT 风格forward 过程是 token 序列形状是 (B, N, D)。这里的 N 是 patch 数量等于feat_h * feat_wD 是每个 patch 的嵌入维度也就是d_model。所以两条分支作用的对象并不一样CBAM 分支作用维度在 TransUnet token 序列里的对应ChannelGateC对 D 维做全局池化即把 token 序列的嵌入维当作通道维SpatialGateH×W把序列 reshape 回特征图在feat_h × feat_w上做空间加权如果只加通道注意力在 (B, N, D) 上直接对 N 做池化就行不需要 reshape。要加完整 CBAMtoken 序列必须能无损 reshape 回 (B, D, feat_h, feat_w)。这个条件看起来简单但很多 ViT 风格实现并不保存分裂后的高宽替换时第一步要确认的就是这一点。reshape 不匹配会直接报错报错信息一般会提示 token 数与期望的feat_h * feat_w不一致。2.2 用不依赖第三方库的 CBAM 实现为改造打底下面这段代码是修改后的基础可以直接复制到模型文件里。为了后面改造 Transformer 块时好调参我把两个门控拆成独立子模块每个子模块都对外暴露方便单独观察通道权重和空间权重的分布。import torch import torch.nn as nn class ChannelGate(nn.Module): def __init__(self, channels, reduction8): super().__init__() hidden max(8, channels // reduction) self.mlp nn.Sequential( nn.Conv2d(channels, hidden, 1, biasFalse), nn.ReLU(inplaceTrue), nn.Conv2d(hidden, channels, 1, biasFalse), ) self.sigmoid nn.Sigmoid() def forward(self, x): avg self.mlp(x.mean(dim(2, 3), keepdimTrue)) max_pool self.mlp(x.amax(dim(2, 3), keepdimTrue)) return self.sigmoid(avg max_pool) class SpatialGate(nn.Module): def __init__(self, kernel_size7): super().__init__() self.conv nn.Conv2d(2, 1, kernel_size, paddingkernel_size // 2, biasFalse) self.sigmoid nn.Sigmoid() def forward(self, x): avg x.mean(dim1, keepdimTrue) max_pool x.amax(dim1, keepdimTrue) return self.sigmoid(self.conv(torch.cat([avg, max_pool], dim1))) class CBAM(nn.Module): def __init__(self, channels, reduction8, kernel_size7): super().__init__() self.ca ChannelGate(channels, reduction) self.sa SpatialGate(kernel_size) def forward(self, x): x x * self.ca(x) x x * self.sa(x) return x代码里有几个点需要说明。平均池化和最大池化是两个不同视图平均池化对整体亮度敏感最大池化更容易捕捉细线条这类局部强响应两者共享 MLP最终逐元素相加再经过 sigmoid得到 0 到 1 之间的通道权重。amax对应 PyTorch 的torch.amax不要写成torch.max后者在二元场景会返回(values, indices)元组直接参与运算会报错。空间注意力把通道维压缩成 2 个特征图7×7 卷积的感受野足够覆盖大多数情况下的局部上下文但特征图尺寸很小时反而会引入噪声后面训练时可以把kernel_size调到 3。这里还隐藏着一个对车道线任务很重要的性质sigmoid 把空间权重限制在 0 到 1 之间意味着这个模块只能做抑制不能做增强。对整幅图不到 2% 像素的前景来说纯抑制有时会连带把细车道线权重压掉。我会在第 5 章给一个改进形式先按标准实现走。2.3 插入位置选在 MSA 之后而不是 MLP 之后Transformer 块里可以插 CBAM 的位置有两处多头自注意力MSA后面或者前馈网络MLP后面。我的建议是放在 MSA 之后用残差形式接进去。原因是 MSA 负责建模全局依赖输出包含大量跨 patch 交互信息CBAM 接在这里能立即对全局响应做一次空间和通道层面的重新校准。如果放在 MLP 后面FFN 已经对特征做了逐 token 线性变换CBAM 再作用相当于二次加权容易把 FFN 里的噪声一并放大。常见做法是写成x x attn(layer_norm(x))之后再把 CBAM 的输出以短路形式加回来。需要注意不要写成x x attn_out cbam_out这种直接拼接残差路径的形式那会造成注意力路径和 CBAM 路径各自独立回传梯度训练早期容易震荡。更稳的写法是先做一次 LayerNorm 再与原始 x 相加下一节直接给出可替换的块实现。3. 重构 TransUnet 的 Transformer 块把 CBAM 接到多头注意力输出上3.1 一个可直接替换的 Transformer 块实现改造的思路把上一节的 CBAM 封装成一个 token 混合器接收 (B, N, D) 序列内部先 reshape 成 (B, D, feat_h, feat_w)应用 CBAM 后再 reshape 回序列。为了让整个模块可以被原版 TransUnet 直接替换对外接口保持(B, N, D) - (B, N, D)。import torch import torch.nn as nn class TransformerBlockWithCBAM(nn.Module): def __init__( self, d_model, nhead, dim_feedforward, dropout0.1, feat_hNone, feat_wNone, reduction8, sp_kernel7, use_cbamTrue, use_caTrue, use_saTrue, ): super().__init__() self.norm1 nn.LayerNorm(d_model) self.attn nn.MultiheadAttention( d_model, nhead, dropoutdropout, batch_firstTrue) self.norm2 nn.LayerNorm(d_model) self.mlp nn.Sequential( nn.Linear(d_model, dim_feedforward), nn.GELU(), nn.Dropout(dropout), nn.Linear(dim_feedforward, d_model), nn.Dropout(dropout), ) self.use_cbam use_cbam self.use_ca use_ca self.use_sa use_sa self.feat_h feat_h self.feat_w feat_w if use_cbam: self.ca ChannelGate(d_model, reduction) if use_ca else nn.Identity() self.sa SpatialGate(sp_kernel) if use_sa else nn.Identity() self.cbam_norm nn.LayerNorm(d_model) def _apply_cbam(self, tokens): b, n, d tokens.shape if n ! self.feat_h * self.feat_w: raise ValueError( ftoken 数 {n} 不匹配 feat_h*feat_w {self.feat_h}*{self.feat_w}) feat tokens.transpose(1, 2).reshape(b, d, self.feat_h, self.feat_w) feat self.ca(feat) feat self.sa(feat) return feat.reshape(b, d, n).transpose(1, 2) def forward(self, x): attn_out self.attn(self.norm1(x), self.norm1(x), self.norm1(x))[0] x x attn_out if self.use_cbam: cbam_out self._apply_cbam(x) x self.cbam_norm(x cbam_out) x x self.mlp(self.norm2(x)) return x这个实现里有几个刻意安排的细节。cbam_norm是对残差相加后的结果做归一化。CBAM 是在图像特征图上计算的输出数值范围与 token 序列的标准分布不同直接往后接 MLP 在训练初期容易出现 loss 不降或者直接发散的情况加一层 LayerNorm 能稳定梯度。另一个细节是把use_ca、use_sa分开消融时不需要改代码初始化模型时传参数即可。注意_apply_cbam里的 reshape 用了浅拷贝视图不会产生额外的显存复制但这也意味着 reshape 前后共享同一块内存任何需要保留下游梯度的操作都应该在 reshape 之前完成。这个实现里先做注意力残差再做 CBAM计算图是连续的不会出现这种情况。3.2 把它接回 TransUnet必须处理的 feat_h 和 feat_w实际替换时feat_h和feat_w是 patch embedding 之后的特征图分辨率。常见做法是直接用卷积 patch embedding那么feat_h img_h // patch_sizefeat_w img_w // patch_size。比如输入 224×224、patch_size16就是 14×14token 数 196。patch_embed nn.Conv2d( in_channels3, out_channelsd_model, kernel_sizepatch_size, stridepatch_size, biasFalse, ) feat_h img_h // patch_size feat_w img_w // patch_size blocks nn.ModuleList([ TransformerBlockWithCBAM( d_modeld_model, nheadnhead, dim_feedforwarddim_feedforward, feat_hfeat_h, feat_wfeat_w, ) for _ in range(num_layers) ])这一段说明了为何feat_h和feat_w必须显式传入。原版 ViT 风格实现里通常只保留 token 序列不关心 N 能不能分解成规则矩形但 CBAM 的空间注意力依赖二维空间结构。如果不传就得在所有 block 内部动态算int(sqrt(N))对于非正方形输入会有明显误差。还有一种常见做法是把feat_h和feat_w作为全局变量在编码器初始化时传入这样所有 block 共用一份计算不容易出错。如果模型里有多级下采样或者 patch merge每个 stage 的feat_h和feat_w都不同不能只传一次。TransUnet 一般不涉及这种结构但改动其他 Transformer 分割模型时要注意。3.3 替换后的参数表与消融建议下面这张表是我在类似分割任务上比较常用的调参依据不是严格的搜索范围但足够作为起步参数。参数推荐值调节说明reduction8通道注意力瓶颈系数数据量小可以调到 16防止过拟合sp_kernel7特征图 14×14 以上用 7降到 7×7 左右建议改 3dropout0.1训练集只有几千张时调 0.2超过两万张保持 0.1cbam_norm保留去掉后训练初期 loss 可能出现平台期每层都加推荐显存紧张时只在最后 6 层加效果略降但省显存替换之后建议先跑两次消融第一次把use_cbam设为False得到 baseline第二次只开use_ca不开use_sa对比通道注意力单独效果。这样能判断改进究竟来自哪一条分支。如果 two-branch 全开反而下降通常是空间注意力的 7×7 卷积在当前特征图上过拟合先改sp_kernel3再考虑 reduction。4. 把改进后的 TransUnet 用到高速道路车道线语义分割上4.1 车道线数据从标注到 mask高速道路车道线数据集里 TuSimple 比较常用标注文件是一个 JSON核心字段是lanes和h_samples。每个 lane 给出等间距高度h_samples上对应的 x 坐标x 为 -1 表示该高度处车道线不存在。下面代码把这种稀疏点转为稠密 mask并把坐标从原始分辨率缩放到模型输入尺寸。import json import numpy as np import cv2 def tusimple_json_to_mask(path, out_size(512, 256), lane_width8): with open(path) as fp: data json.load(fp) mask np.zeros((out_size[1], out_size[0]), dtypenp.uint8) lanes data[lanes] h_samples data[h_samples] for lane in lanes: pts [] for y, x in zip(h_samples, lane): if x 0: pts.append([int(x), int(y)]) if len(pts) 2: continue x_scale out_size[0] / 1280.0 y_scale out_size[1] / 720.0 pts [[int(p[0] * x_scale), int(p[1] * y_scale)] for p in pts] cv2.polylines(mask, [np.array(pts)], False, 1, thicknesslane_width) return mask这段代码的逻辑是先按原始坐标收集二维点再用cv2.polylines把离散点连成连续线。lane_width控制线条粗细相当于给车道线标注一个先验宽度8 个像素在 512×256 输入下大约覆盖 1.5% 的像素正负样本比更容易控制。如果直接把点画成单个像素绝大多数样本的前景占比会低于 0.5%BCE 损失非常容易把前景全部学成背景。TuSimple 的原始分辨率是 1280×720缩放系数直接按目标宽高计算。如果模型输入是 224×224建议先缩放再裁剪下半部分道路区域不要直接把 1280×720 拉伸到 224×224否则远处车道线的宽窄比例会被扭曲影响空间注意力学习和最终 F1 值。4.2 为少数类设计的损失函数与 F1 评估车道线语义分割的类别严重不平衡纯 BCE 会偏向背景。业界比较稳的组合是 BCE 加 Dice 损失。Dice 损失对前景比例不敏感能增强细长目标上的梯度信号而 BCE 又能提供稠密的像素级监督两者互补。def dice_loss(pred, target, smooth1.0): pred torch.sigmoid(pred) inter (pred * target).sum(dim(1, 2)) union pred.sum(dim(1, 2)) target.sum(dim(1, 2)) return 1 - (2 * inter smooth) / (union smooth) loss F.binary_cross_entropy_with_logits(pred, target) dice_loss(pred, target)dice_loss里的union用 pred 与 target 各自求和再相加实现避免直接算交集和并集时对全零样本除零。全零 mask 的预测也接近全零时union接近 0分母由smooth1.0兜底。评估指标不建议直接用逐像素 IoU因为背景占比太高IoU 会被稀释到看起来不错但车道线实际断成一截的状态。常见做法是像素级 F1预测概率以 0.3 为阈值二值化统计整张验证集的 TP、FP、FN最后汇总计算精确率、召回率和 F1。注意要汇总后再算如果逐样本算 F1 再平均没有车道线的样本会拉低结果。4.3 训练命令与高频踩坑对照表训练时建议显式区分两类输入尺寸不要用同一个参数控制。模型输入 224×224 比较省显存但车道线远端像素损失太多512×256 保留了更多横向信息推荐作为主配置。python train.py \ --data-dir data/tusimple/train_set/ \ --img-size 512 256 \ --patch-size 16 \ --batch-size 8 \ --lr 3e-4 \ --warmup-epochs 2 \ --max-epochs 60 \ --f1-threshold 0.3 \ --save-ckpt ckpt/transunet_cbam/--f1-threshold是二值化阈值。车道线模型的预测概率分布通常右偏0.3 比 0.5 更容易把远处置信度偏低的点连起来。如果验证时发现误检很多可以往 0.5 调如果断裂多往 0.25 调。warmup-epochs2对加了 CBAM 的 Transformer 格外重要因为通道注意力的初始权重会改变优化曲面直接大学习率起步容易崩。下面是这个任务上常见的几个坑以及我一般会采用的排查路径。现象常见原因处理方式loss 降不下去mask 与图像坐标没对齐可视化一张 mask 与原图的叠加图检查 ROI 裁剪加了 CBAM 后 F1 掉点空间注意力把细线响应压掉了开use_caTrue, use_saFalse消融调整为 3×3 kernel验证集断裂多阈值设置偏高--f1-threshold降到 0.25 或 0.3显存涨了 30% 以上每层都做 reshape 和 7×7 卷积只保留最后 6 层 CBAM或把 sp_kernel 改小5. 两个能直接提升效果的落地技巧5.1 先验证 CBAM 板块真的拿到了梯度很多人改完模型直接训练发现效果不明显就否定了方案。但最常见的情况是 CBAM 分支在某个 forward 分支里根本没生效。换模型结构之后第一步不是看 F1而是检查梯度通路。model.train() pred model(images) loss criterion(pred, masks) loss.backward() block model.encoder.blocks[3] if block.use_cbam: grad_ca block.ca.mlp[0].weight.grad grad_sa block.sa.conv.weight.grad print(ChannelGate grad:, None if grad_ca is None else grad_ca.abs().mean().item()) print(SpatialGate grad:, None if grad_sa is None else grad_sa.abs().mean().item())如果两个梯度都是None说明该 block 的 CBAM 根本没参与计算。检查点通常有两个一个是_apply_cbam里的 reshape 是否在 forward 中被跳过另一个是 Python 的if self.use_cbam分支在初始化后有没有被父模块误改。如果梯度不为零但数值在 1e-8 量级说明注意力权重对 loss 的贡献微乎其微大概率是通道维被 MLP 压到 1 附近sigmoid 输出饱和这种情况可以调大reduction值重新初始化。5.2 给空间注意力做残差化处理标准 CBAM 的 spatial gate 输出范围是 0 到 1作用只能是衰减。但车道线检测里远处线条原本响应就弱衰减过多会导致断线。常用的改进是给空间权重叠加一个恒等偏置让模块既可以抑制背景也能增强前景。class SpatialGateV2(nn.Module): def __init__(self, kernel_size3, init_bias1.0): super().__init__() self.conv nn.Conv2d(2, 1, kernel_size, paddingkernel_size // 2) nn.init.constant_(self.conv.bias, init_bias) def forward(self, x): avg x.mean(dim1, keepdimTrue) max_pool x.amax(dim1, keepdimTrue) w torch.sigmoid(self.conv(torch.cat([avg, max_pool], dim1))) return x * (1 w)做法是在原特征上乘以(1 sigmoid(...))权重范围变成 1 到 2模块可以增强任意空间位置同时仍然保持对背景的相对抑制。init_bias1.0让初始权重在 2 左右开始训练比较稳定。把 kernel_size 从 7 改为 3是考虑到 TransUnet 的 token 特征图通常只有 14×14 或 7×77×7 卷积几乎覆盖全图空间注意力退化成全局亮度调整失去局部定位能力。用这个替换标准SpatialGate后建议在场景多变的验证集上同时看 F1 和误检率避免增强后把路肩、护栏等边缘也一并拉高。本文还有配套的精品资源点击获取