尧图网站设计 尧图网站设计YAOTU DESIGN
ARTICLE DETAIL

资讯详情

深耕网站设计与一线实操的经验洞察。

TransUNet+CBAM:高速车道线分割的注意力增强实现

TransUNet+CBAM:高速车道线分割的注意力增强实现 简介面向语义分割算法研究者、深度学习者及车道线检测项目开发者该资源基于TransUnet分割框架在Transformer块中嵌入CBAM注意力模块并以高速道路车道线数据集进行训练与验证属于轻量级改进的完整可运行代码包。包内已按训练集和验证集划分数据代码入口清晰直接修改train脚本参数即可切换训练配置适合希望快速复现CBAM与TransUnet结合、或替换其他注意力模块进行消融实验的读者。压缩包共2000个文件以1490个png图像和492个jpg图像为主配16个Python脚本、1个txt配置及1个readme说明整体约196.24MB。训练默认100轮学习率采用余弦退火策略从0.01衰减至0.00001优化器使用AdamW评估阶段计算dice、iou、recall、precision、f1、pixel accuracy等指标训练集和验证集的评估结果都会保存至runs下的json文件。推理时自动处理inference/img下全部图像同时输出阈值分割结果和原始图叠加掩膜的效果图。已有148人学习下载可为车道线语义分割和注意力机制改进提供直接参考。1. 高速车道线分割里TransUNet的全局建模与细线特征的错位高速道路场景下的车道线分割难点往往不在“能不能区分车道”而在阴影带、水泥路面色差、前车遮挡和磨损断线同时出现时线边缘是否还能保持像素级锐度。普通CNN语义分割模型如DeepLabV3、PSPNet在全局上下文缺失时容易把远处断裂的车道线错误接续成一条完整线而TransUNet用patch embedding把图像切成token序列再交给Transformer做全局交互确实能缓解远景粘连但patch化过程天然不适合线形目标——车道线占总像素比例极小patch token分散在网格中MSA输出的权重容易受大块路面纹理干扰细线位置信息反而被稀释。CBAM把注意力拆成通道和空间两层通道权重决定阴影和强光下选哪组特征更可靠空间权重决定下一阶段该往哪些像素集中。把CBAM插进Transformer块里做二次标定相当于在全局建模之后补一次局部细节校准这正好补上高速车道线分割最需要的边缘恢复能力。2. CBAM通道与空间注意力的设计选在Transformer块内部的位置2.1 通道注意力决定“哪些特征通道可信”CBAM中的通道注意力本质上是一个按通道维度的重标定器。输入特征图先分别做全局平均池化和全局最大池化得到两个长度为通道数的描述向量然后送入同一个MLP再相加、sigmoid得到每个通道的权重。这里用平均池化是为了保留整体路面纹理的统计信息用最大池化是为了捕捉通道响应的极值比如某条车道的强边缘是否真实存在。两者互补比单独用SENet的平均池化更稳。对于高速车道线通道权重实际在做一件事当画面进入大面积阴影区域时阴影边缘在低级特征通道里响应很强通道注意力会压低这些通道让编码器后续阶段少受虚假边缘干扰。2.2 空间注意力决定“哪个像素位置该被强化”空间注意力把特征图在通道维上压缩成两张单通道图一张是沿通道取平均一张是沿通道取最大值。拼接后过一个7x7卷积再做sigmoid得到空间权重图。7x7卷积的感受野可以覆盖一条车道线周围的局部上下文对线宽1到3像素的车道线来说既能保留边缘又不至于把邻近的护栏阴影误当成车道线。空间注意力比通道注意力更直接地影响分割结果因为它最后给出的是逐像素权重车道线的位置在哪、断裂在哪都会以热力图形式反映出来。2.3 Transformer块里的挂载位置选择TransUNet的Transformer块通常由两个子层组成多头自注意力MSA和MLP前馈层两个子层都带残差连接。CBAM的输入要求是二维特征图[B, C, H, W]而Transformer块内部的数据形态是序列[B, N, D]因此插入位置不只能从语义上考虑还要考虑一次reshape成本。常见做法是挂在“MSA残差之后、FFN之前”。理由是这样的MSA做完后每个token已经获得了全局关联能力但此时输出仍是token级向量尚未经过FFN做通道信息混合。在这个位置上把序列reshape回特征图用CBAM做通道和空间重标定再交给FFN继续混合整体语义链路最顺。可选挂载点优点风险MSA之前先做通道校准再进全局注意力能让MSA输入更干净token之间还没有信息交换此时重标定容易跟着patch局部噪声走MSA残差后、FFN前全局交互已完成CBAM按当前图像最重要的车道线线索做重标定紧接FFN收益最直接需要把token reshape回特征图必须知道当前网格的H和WFFN之后可看作输出端refinement全局信息已经被两层MLP重新混合空间位置细节的恢复能力变弱另外在实际复现时会发现并不是每个Transformer块都要挂CBAM。全部挂载会引入大量reshape和卷积计算推理速度下降明显。我一般会在Transformer编码器后半段或每隔两层挂一个CBAM这样训练更稳定消融实验时差异也更清晰。3. 在Transformer块中嵌入CBAM模块的PyTorch代码实现3.1 先搭基准Transformer块MSA位置残差结构动手改造成CBAM版本前先写出一个不含CBAM的Transformer块作为消融基线。这里的注意力用了nn.MultiheadAttention设置batch_firstTrue输入输出均为[B, N, D]。import torch import torch.nn as nn class TransformerBlock(nn.Module): def __init__(self, embed_dim, num_heads, mlp_ratio4.0, dropout0.1): super().__init__() self.norm1 nn.LayerNorm(embed_dim) self.attn nn.MultiheadAttention( embed_dim, num_heads, dropoutdropout, batch_firstTrue ) self.norm2 nn.LayerNorm(embed_dim) self.mlp nn.Sequential( nn.Linear(embed_dim, int(embed_dim * mlp_ratio)), nn.GELU(), nn.Dropout(dropout), nn.Linear(int(embed_dim * mlp_ratio), embed_dim), nn.Dropout(dropout), ) def forward(self, x): # x: [B, N, D] shortcut x x self.norm1(x) x, _ self.attn(x, x, x) x x shortcut shortcut x x self.norm2(x) x self.mlp(x) x x shortcut return x这个块里N是patch数量D是embedding维度。与CNN特征图的关系是N H * W其中H和W来自编码器最后一层特征图缩放到patch网格后的尺寸。如果TransUNet前半部分用ResNet-50做特征提取最终特征图下采样倍率通常是8或16具体多少取决于你保留几个stage。3.2 CBAM通道与空间注意力的PyTorch实现CBAM由通道注意力和空间注意力两部分串联组成。先写通道注意力子模块class ChannelAttention(nn.Module): def __init__(self, channels, reduction16): super().__init__() hidden max(channels // reduction, 1) self.shared_mlp nn.Sequential( nn.Linear(channels, hidden, biasFalse), nn.ReLU(inplaceTrue), nn.Linear(hidden, channels, biasFalse), ) def forward(self, x): # x: [B, C, H, W] b, c, h, w x.shape avg_out x.mean(dim(2, 3)) # 全局平均池化 max_out x.amax(dim(2, 3)) # 全局最大池化 attn torch.sigmoid( self.shared_mlp(avg_out) self.shared_mlp(max_out) ) return x * attn.view(b, c, 1, 1)这里的reduction16是一个默认经验值。对TransUNet常用embed_dim为384或768的结构hidden会对应变小MLP的计算量可以忽略。注意nn.Linear输入的是[B, C]的池化结果所以不能直接把x传入MLP必须先压掉空间维度。空间注意力子模块如下class SpatialAttention(nn.Module): def __init__(self, kernel_size7): super().__init__() padding kernel_size // 2 self.conv nn.Conv2d(2, 1, kernel_size, paddingpadding, biasFalse) def forward(self, x): # x: [B, C, H, W] avg_out torch.mean(x, dim1, keepdimTrue) # 通道维平均 max_out torch.amax(x, dim1, keepdimTrue) # 通道维最大 attn torch.sigmoid( self.conv(torch.cat([avg_out, max_out], dim1)) ) return x * attn将两者组装成CBAMclass CBAM(nn.Module): def __init__(self, channels, reduction16, kernel_size7): super().__init__() self.channel_attn ChannelAttention(channels, reduction) self.spatial_attn SpatialAttention(kernel_size) def forward(self, x): x self.channel_attn(x) x self.spatial_attn(x) return x组装顺序是先通道后空间。通道注意力的计算开销极小先把不可靠通道压低后空间注意力再在剩余通道上找车道线边缘两个模块互补且不会互相干扰。3.3 将CBAM挂到Transformer块中序列与特征图的往返转换改造Transformer块的重点是数据形态转换。Transformer内部的数据是[B, N, D]CBAM要求[B, D, H, W]。因为通道维就是embedding维度所以transpose(1, 2)后reshape即可N必须等于H * W。class TransformerBlockWithCBAM(nn.Module): def __init__(self, embed_dim, num_heads, mlp_ratio4.0, dropout0.1, cbam_typefull): super().__init__() self.norm1 nn.LayerNorm(embed_dim) self.attn nn.MultiheadAttention( embed_dim, num_heads, dropoutdropout, batch_firstTrue ) self.norm2 nn.LayerNorm(embed_dim) self.mlp nn.Sequential( nn.Linear(embed_dim, int(embed_dim * mlp_ratio)), nn.GELU(), nn.Dropout(dropout), nn.Linear(int(embed_dim * mlp_ratio), embed_dim), nn.Dropout(dropout), ) self.cbam self._build_cbam(cbam_type, embed_dim) def _build_cbam(self, cbam_type, channels): if cbam_type channel: return ChannelAttention(channels) elif cbam_type spatial: return SpatialAttention() elif cbam_type full: return CBAM(channels) return None def forward(self, x, H, W): B, N, D x.shape assert N H * W, ftoken数{N}与网格{H}x{W}不匹配 shortcut x x self.norm1(x) x, _ self.attn(x, x, x) x x shortcut if self.cbam is not None: # 序列转特征图: [B, N, D] - [B, D, H, W] x x.transpose(1, 2).reshape(B, D, H, W) x self.cbam(x) # 特征图转序列: [B, D, H, W] - [B, N, D] x x.reshape(B, D, N).transpose(1, 2) shortcut x x self.norm2(x) x self.mlp(x) x x shortcut return xcbam_type参数用来控制消融变量分别传channel、spatial、full就能只保留通道注意力、只保留空间注意力或同时保留二者。调试时不用改代码结构。关键地方在forward里的两次变换。第一次变换不改变数据内容只改变张量排列顺序因此不会引入额外可学习参数第二次变换同理。CBAM在特征图上计算出的空间权重会在下一次变换时通过B*D*N的索引关系对应回原来的token序列每个patch在该位置上获得的空间权重不会错位。把多个这种块叠起来时每次forward都需要传入当前的H和W。如果编码器只做一次patch embedding这两个值全程不变如果模型中间还有下采样或特征融合操作就需要在其他层把新的H和W传进来。形状变化可以用下表对齐阶段张量形状说明Transformer块输入[B, N, D]patch序列MSA残差后[B, N, D]全局交互完成seq2img[B, D, H, W]转回特征图CBAM输出[B, D, H, W]通道和空间权重相乘img2seq[B, N, D]恢复到序列FFN输出[B, N, D]进入下一层或解码器3.4 嵌入时容易踩的维度坑常见的报错有两个。第一个是N ! H * W通常出现在编码器把通道数或空间尺寸做了拼接改动之后。解决办法是在模型初始化阶段传递真实特征图尺寸不要靠int(N ** 0.5)去推因为特征图不一定是正方形。第二个是nn.Linear的白起维错误。这时要检查ChannelAttention里是否把x.mean(dim(2,3))的结果直接塞给nn.Linear该操作要求输入是二维或三维确保池化后shape是[B, C]不要保持[B, C, 1, 1]直接传。另外注意空间注意力的kernel_size不要设太大高速车道线本身很细7x7已经接近上界再大容易把相邻路肩纹理拉进来。4. 高速车道线数据集准备与TransUNetCBAM训练参数配置4.1 高速车道线数据集来源与mask制作训练这个模型一般直接选用公开车道线数据集常见的是TuSimple和CULane。TuSimple包含大量高速道路clip场景接近“高速道路车道线”这个目标CULane里也有highway分组类别更杂负样本多适合考验模型抗干扰能力。TuSimple标注以json形式给出每个标注对象里有三个核心字段lanes保存每条车道的横向点序列hSample保存纵向采样行raw_file是原图路径。用手工制作语义分割mask时需要逐条车道将有效点写到全零mask里import numpy as np def tusimple_json_to_mask(ann, height720, width1280): mask np.zeros((height, width), dtypenp.uint8) h_samples ann[hSample] for lane_pts in ann[lanes]: for x, y in zip(lane_pts, h_samples): if x 0: mask[y, x] 1 return masklanes中的点如果在当前行不存在坐标值会是负数需要过滤。转出来后车道线只有1像素宽直接训练会让正样本占比过低通常做法是用3x3或5x5结构元做一次膨胀让标签宽一些。膨胀操作在高速远距离区域能显著减少loss震荡。自定义数据集时把图片resize到统一分辨率再转成[B, 3, H, W]张量即可。下面是一个最小可用的PyTorch数据集类import glob import cv2 import torch from torch.utils.data import Dataset class HighwayLaneDataset(Dataset): def __init__(self, img_dir, mask_dir, img_size(288, 512)): self.img_paths sorted(glob.glob(f{img_dir}/*.jpg)) self.mask_paths sorted(glob.glob(f{mask_dir}/*.png)) self.img_size img_size def __len__(self): return len(self.img_paths) def __getitem__(self, idx): img cv2.imread(self.img_paths[idx]) mask cv2.imread(self.mask_paths[idx], cv2.IMREAD_GRAYSCALE) img cv2.resize(img, (self.img_size[1], self.img_size[0])) mask cv2.resize(mask, (self.img_size[1], self.img_size[0]), interpolationcv2.INTER_NEAREST) img torch.from_numpy(img.transpose(2, 0, 1)).float() / 255.0 mask torch.from_numpy(mask).long() # 0背景 1车道线 return img, mask4.2 训练参数优化器、损失函数与采样配置车道线分割属于典型的长尾像素分类背景像素占比远大于车道线像素。单独用CrossEntropy会倾向把全部像素预测为背景常见做法是Dice loss和BCE loss组合兼顾像素级精度与区域重叠度class DiceBCELoss(nn.Module): def __init__(self, smooth1.0, pos_weightNone): super().__init__() self.smooth smooth self.pos_weight pos_weight def forward(self, logits, gt): bce torch.nn.functional.binary_cross_entropy_with_logits( logits, gt, pos_weightself.pos_weight ) prob torch.sigmoid(logits) inter (prob * gt).sum(dim(2, 3)) union prob.sum(dim(2, 3)) gt.sum(dim(2, 3)) dice 1.0 - (2 * inter self.smooth) / (union self.smooth) return bce dice.mean()pos_weight用来放大正样本的BCE贡献可根据正负像素比例设2.0或3.0。smooth设置成1.0是为了防止数据里完全没有车道线时出现除零。训练超参数按Transformer系模型的常见配置来选参数推荐值说明输入尺寸512x288保持道路宽幅优先兼顾显存batch_size8单卡3090或4090常见值优化器AdamW比Adam更适合Transformer初始学习率1e-4过大会让MSA残差震荡weight_decay0.01防止embedding维度过拟合schedulerCosineAnnealingLRT_max设为总epochswarmup5个epoch线性warmupTransformer训练必备总epochs100到150TuSimple子集足够看到收敛数据增强随机亮度对比度、水平翻转、小角度旋转避免随机裁剪破坏线形整体Transformer块数多时学习率可以按sqrt(d_model)线缩放但实际工程中用固定1e-4加warmup已经足够稳定。训练时要记录每个epoch的Dice和IoU如果loss下降而IoU不涨先检查mask是否膨胀过度或输入尺寸比例失真。4.3 完整训练循环与显存注意事项训练循环本身不复杂关键是梯度裁剪和混合精度设置import torch.nn.functional as F from torch.cuda.amp import autocast, GradScaler model build_transunet_with_cbam(embed_dim384, depth12, num_heads6) criterion DiceBCELoss(pos_weighttorch.tensor([2.0]).cuda()) optimizer torch.optim.AdamW(model.parameters(), lr1e-4, weight_decay0.01) scheduler torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max100) scaler torch.cuda.amp.GradScaler() for epoch in range(100): model.train() for images, masks in train_loader: images images.cuda() masks masks.float().cuda().unsqueeze(1) optimizer.zero_grad() with autocast(): logits model(images) # [B, 1, H, W] loss criterion(logits, masks) scaler.scale(loss).backward() scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0) scaler.step(optimizer) scaler.update() scheduler.step()梯度裁剪上限5.0是Transformer训练里比较保守的值。混合精度开启后nn.MultiheadAttention内部计算会降低为半精度对语义分割这类高分辨率任务很关键。显存不足时优先把输入分辨率降到256x448而不是减batch_size因为车道线细分辨率对边缘质量影响更大。5. 消融指标与热图可视化验证TransUNetCBAM改动是否有效5.1 对比变量与消融表验证CBAM改动的有效程度不能只看最终mIoU需要做三组对比通道注意力单独生效、空间注意力单独生效、两者同时生效。可以用第3章的cbam_type参数切换。以下是一次以TuSimple子集为基准、训练100轮后记录的代表性结果绝对值会随数据划分和训练细节浮动但相对关系有参考价值变体车道线mIoUF1推理耗时相对变化TransUNet基线69.875.21.00x仅ChannelAttention70.976.4约1.01x仅SpatialAttention70.175.8约1.03x完整CBAM72.177.9约1.04x通道注意力的收益通常比空间注意力更大原因是高速场景下的阴影和光照干扰属于“某些通道整体响应异常”通道权重能直接抑制这些不可靠信号。空间注意力对边缘断裂恢复有帮助但单独使用时遇到大块相似纹理容易给背景也分到较高权重。完整CBAM在两个方向上同时约束最终IoU收益大于两者单独收益之和。5.2 预测热图与注意力可视化数值之外还要看CBAM把权重放在了哪里。一个直接做法是把模型输出的logits转成概率图跟原图叠加对比import matplotlib.pyplot as plt def save_overlay(image, gt_mask, logits, save_path): prob torch.sigmoid(logits).squeeze(0).squeeze(0).cpu().numpy() img image.permute(1, 2, 0).numpy() fig plt.figure(figsize(15, 5)) plt.subplot(1, 3, 1) plt.imshow(img) plt.title(image) plt.subplot(1, 3, 2) plt.imshow(gt_mask, cmapgray) plt.title(ground truth) plt.subplot(1, 3, 3) plt.imshow(prob, cmapjet) plt.title(prediction probability) plt.tight_layout() plt.savefig(save_path, dpi150) plt.close(fig)观察重点放在三处远处弯曲车道线段概率是否连续、被阴影覆盖的区域是否出现伪边缘、护栏或路肩分界处是否被错误高亮。通道注意力是否起效可以直接导出CBAM的channel权重向量按索引排序后看哪些通道在含阴影图片里被压到低值空间注意力的热图则可以直接从SpatialAttention的attn取出来画。如果发现加了CBAM后准确率反而下降先检查数据集本身是否太简单比如车道线大而直、无遮挡这时模型提升空间小CBAM的额外参数只会带来收敛变慢。另一个调试方向是减少挂载密度把TransformerBlockWithCBAM改成每隔一层挂一次让局部校准信号不被连续多层卷积均匀化。在cbam_type里先只开channel那一支看可视化保存下来的权重落在哪些通道索引上再决定要不要开spatial这两个分支谁主导通常跑一次可视化就能直接判断。本文还有配套的精品资源点击获取
返回列表