1. 从“看局部”到“看全局”为什么语义分割需要双注意力机制如果你玩过“大家来找茬”或者“找不同”的游戏你肯定有过这样的体验有时候两个地方的差异非常细微你盯着一个局部看了半天也发现不了但当你退后一步把整张图都扫一眼对比一下不同区域的整体感觉那个不同点可能一下子就跳出来了。传统的卷积神经网络CNN在做语义分割——也就是给图片里的每个像素都打上标签分出哪里是天空、哪里是道路、哪里是人——的时候就有点像那个只盯着局部看的人。CNN通过一层层的卷积操作确实能提取出很棒的特征比如边缘、纹理。但它的“感受野”是有限的。你可以把感受野理解成网络在某一层“看”到的原始图像区域大小。即使经过很多层这个范围可能还是不够大。这就导致了一个问题对于一张大图里两个相隔很远但属于同一类别的物体比如图片左上角的一辆车和右下角的另一辆车网络在分别处理它们所在的局部区域时可能因为感受野的局限无法建立它们之间的联系从而可能把它们识别成不同的东西或者分割得不够连贯。这就像你只看一棵树的树干很难判断它是松树还是杨树但如果你能看到整棵树的树冠形状判断起来就准确多了。我们需要一种机制让网络在分析某个像素时不仅能参考它周围那一小片区域的特征还能“回忆”或“注意到”整张图片里所有和它相似的特征区域。这就是注意力机制的核心思想让模型学会“聚焦”在重要的信息上并建立长距离的依赖关系。而DAnetDual Attention Network的聪明之处在于它觉得只从一个维度“聚焦”还不够过瘾。它引入了双注意力机制相当于给模型配了两副功能不同的“眼镜”空间注意力Position Attention这副眼镜帮助模型在二维平面空间上建立联系。无论两个像素点隔得多远只要它们的特征相似比如都是“车”的金属质感这副眼镜就能把它们关联起来让它们的预测结果相互加强、保持一致。通道注意力Channel Attention这副眼镜帮助模型在特征通道维度上做筛选。我们都知道卷积后会产生很多个特征通道Channel每个通道可以粗略地理解为对某种特定模式比如横向边缘、红色区域、车轮纹理的“探测器”。通道注意力机制的作用就是自动判断在当前场景下哪些通道的“探测结果”更重要然后给它们“调高音量”抑制那些不重要的通道。我打个比方你要在一张复杂的街景图中找出所有的行人。空间注意力就像让你同时扫视图片的各个角落把所有看起来像“人形”的局部区域都关联起来不管他们在前景还是背景。而通道注意力则像是一个经验丰富的助手他会告诉你“在这个场景里关注‘肤色’、‘衣物纹理’和‘直立轮廓’这几个特征通道特别有用至于‘天空蓝’和‘柏油路灰’这些通道暂时可以放一放。”双管齐下模型既能从空间上理解物体的完整分布和上下文关系又能从特征本质通道上抓住最具有判别性的信息。这就是DAnet在CVPR 2019上提出后效果显著的原因。它没有使用特别复杂或昂贵的结构比如空洞空间金字塔池化ASPP那种多分支并行而是用相对优雅的矩阵运算实现了强大的全局上下文建模能力。接下来我们就亲手用PyTorch把它实现出来看看这两副“眼镜”到底是怎么工作的。2. 庖丁解牛双注意力模块的PyTorch实现详解理解了为什么需要双注意力我们直接来看代码怎么写。这是整个DAnet最核心、最精彩的部分。我会把论文里的公式和图表翻译成你能看懂的操作步骤和代码行。2.1 空间注意力模块让像素们“对话”空间注意力模块的目标是生成一个N x N的注意力图N 高度 H * 宽度 W。这个图的每一个元素S[j, i]都代表了位置i对位置j的重要性权重。换句话说它描述了图片中任意两个像素点之间的“亲密度”。它是怎么算出来的呢我们一步步拆解特征变换输入一个特征图A形状是(批大小B, 通道数C, 高H, 宽W)。我们通过三个独立的1x1卷积层把它变换成三个新的特征图B,C,D。注意这里的C既是通道数也是第二个特征图的变量名别搞混了。1x1卷积的作用是进行通道间的信息融合和降维如果需要的话在这里主要是为了生成用于计算相似度的投影。计算相似度矩阵把B和C的形状从(B, C, H, W)变成(B, C, N)其中N H*W。然后进行矩阵乘法B^T * C在代码中体现为对B进行转置后与C相乘。这个操作厉害了它让B的每一个空间位置共N个的特征向量都与C的每一个空间位置的特征向量计算了点积相似度。结果就是一个(B, N, N)的矩阵它包含了所有像素对之间的原始关联度。归一化得到注意力图对这个(B, N, N)的矩阵在最后一个维度dim2上应用softmax归一化。这样对于目标位置j对应矩阵的第j行所有源位置i对应矩阵的列对其的权重之和就为1。权重越大说明位置i的特征对位置j越重要。加权聚合用得到的注意力图S去加权聚合特征。将特征图D也reshape成(B, C, N)然后与注意力图S的转置进行矩阵乘法D * S^T。这个操作意味着对于每一个位置j的新特征是所有位置i的特征D[:, :, i]以S[j, i]为权重加权求和的结果。这就把全局的上下文信息聚合到了每个位置上。残差连接最后将聚合后的特征图reshape回(B, C, H, W)乘以一个可学习的权重参数gamma初始化为0再加上最原始的特征图A。残差连接是深度学习里的经典技巧它确保了加入注意力模块后网络至少不会比原来更差让训练更稳定。随着训练gamma会逐渐学习到一个合适的值决定依赖多少全局上下文信息。下面是我根据论文复现的PositionAttention模块的PyTorch代码加上了详细的注释import torch import torch.nn as nn class PositionAttention(nn.Module): def __init__(self, in_channels): super(PositionAttention, self).__init__() # 三个1x1卷积分别生成B, C, D特征图 self.convB nn.Conv2d(in_channels, in_channels, kernel_size1, biasFalse) self.convC nn.Conv2d(in_channels, in_channels, kernel_size1, biasFalse) self.convD nn.Conv2d(in_channels, in_channels, kernel_size1, biasFalse) # 可学习的权重参数gamma初始化为0通过残差连接逐渐学习 self.gamma nn.Parameter(torch.zeros(1)) self.softmax nn.Softmax(dim2) # 在特征维度上做softmax def forward(self, x): 输入 x: (B, C, H, W) 输出: (B, C, H, W) 增强了空间上下文信息的特征图 batch_size, c, h, w x.size() N h * w # 像素总数 # 步骤1: 生成B, C, D feat_B self.convB(x) # (B, C, H, W) feat_C self.convC(x) # (B, C, H, W) feat_D self.convD(x) # (B, C, H, W) # 步骤2: reshape并计算相似度 # 将B, C reshape成 (B, C, N) 然后转置B为 (B, N, C) 以便与C (B, C, N) 做矩阵乘 feat_B_reshaped feat_B.view(batch_size, c, N) # (B, C, N) feat_C_reshaped feat_C.view(batch_size, c, N) # (B, C, N) # 计算相似度矩阵: (B, N, C) (B, C, N) - (B, N, N) # 这个矩阵的 [j, i] 元素表示位置i对位置j的影响权重 attention_s torch.bmm(feat_B_reshaped.transpose(1, 2), feat_C_reshaped) # (B, N, N) # 步骤3: 归一化得到空间注意力图 attention_s self.softmax(attention_s) # (B, N, N) # 步骤4: 用注意力图加权聚合特征D feat_D_reshaped feat_D.view(batch_size, c, N) # (B, C, N) # (B, C, N) (B, N, N) - (B, C, N) 注意这里用的是attention_s的转置 # 实际上我们通常用 S^T 与 D 乘论文图示和公式可能因视角不同有差异但本质都是加权求和。 # 更常见的写法是 context torch.bmm(feat_D_reshaped, attention_s.transpose(1, 2)) context torch.bmm(feat_D_reshaped, attention_s.transpose(1, 2)) # (B, C, N) context context.view(batch_size, c, h, w) # (B, C, H, W) # 步骤5: 残差连接 output self.gamma * context x return output注意这里有一个容易混淆的点。有些实现和论文图示中计算加权聚合时使用的是注意力图S本身有些使用的是S的转置。这取决于你如何定义“影响”。如果你认为S[j,i]是i对j的影响那么用D乘以S^T来为每个j聚合信息就是正确的因为矩阵乘法是行乘列。关键是要理解其物理意义每个位置的新特征是所有位置旧特征的加权和权重由它们之间的相似度决定。2.2 通道注意力模块给特征通道“打分”通道注意力模块的结构与空间注意力类似但操作对象从“空间位置”变成了“特征通道”。它的目标是生成一个C x C的注意力图X其中X[j, i]表示通道i对通道j的重要性。它的计算流程如下直接利用原始特征与空间注意力不同通道注意力模块通常不对输入特征A做额外的卷积变换。作者认为这样可以保持原始通道间的关系。所以我们直接使用A。计算通道间相似度将A从(B, C, H, W)reshape成(B, C, N)。然后进行矩阵乘法A * A^T。这实际上是计算了不同通道的特征图之间的相关性。如果两个通道的特征图总是同时激活比如“车轮”通道和“车身”通道它们的相关性就会很高。归一化得到通道注意力图对得到的(B, C, C)矩阵在最后一个维度应用softmax得到通道注意力图X。加权聚合与残差连接将注意力图X与reshape后的特征A相乘再reshape回原始形状最后通过一个可学习参数beta进行残差连接。class ChannelAttention(nn.Module): def __init__(self): super(ChannelAttention, self).__init__() # 可学习的权重参数beta初始化为0 self.beta nn.Parameter(torch.zeros(1)) self.softmax nn.Softmax(dim2) def forward(self, x): 输入 x: (B, C, H, W) 输出: (B, C, H, W) 增强了通道间依赖关系的特征图 batch_size, c, h, w x.size() N h * w # 步骤12: reshape并计算通道间相似度 feat_reshaped x.view(batch_size, c, N) # (B, C, N) # 计算通道注意力图: (B, C, N) (B, N, C) - (B, C, C) attention_x torch.bmm(feat_reshaped, feat_reshaped.transpose(1, 2)) # (B, C, C) # 步骤3: 归一化 attention_x self.softmax(attention_x) # (B, C, C) # 步骤4: 加权聚合与残差连接 # (B, C, C) (B, C, N) - (B, C, N) context torch.bmm(attention_x, feat_reshaped) # (B, C, N) context context.view(batch_size, c, h, w) # (B, C, H, W) output self.beta * context x return output这个模块理解起来可能更直观一些。假设我们的特征有256个通道这个模块就会计算出一个256x256的矩阵。这个矩阵相当于一个“通道关系表”它告诉网络“在当前这个图像区域里如果你想加强第50个通道比如它对应‘玻璃’特征的响应那么你应该同时去参考第120通道‘窗户框’特征和第200通道‘反光’特征的信息因为它们经常同时出现。”2.3 双注意力头DAHead的组装有了两个核心模块我们需要把它们组装到一个解码器头部Head里。DAnet的头部结构也很清晰将主干网络提取的深层特征例如来自ResNet-50的2048维特征通过两个独立的卷积分支进行预处理得到两个特征图分别送入空间注意力模块和通道注意力模块。论文中这里会先进行降维例如降到256维。将两个注意力模块的输出直接相加element-wise sum。这里蕴含的思想是空间和通道信息是互补的简单相加就能融合。对融合后的特征再进行几层卷积逐步上采样并恢复到原始图像大小同时将通道数映射到类别数得到最终的像素级分类预测图。class DAHead(nn.Module): def __init__(self, in_channels, num_classes, reduction4): in_channels: 主干网络输出通道数如ResNet50的2048 num_classes: 分割类别数 reduction: 降维倍数默认为4即中间特征通道数 in_channels // reduction super(DAHead, self).__init__() mid_channels in_channels // reduction # 两个预处理卷积分支 self.conv_pa nn.Sequential( nn.Conv2d(in_channels, mid_channels, kernel_size3, padding1, biasFalse), nn.BatchNorm2d(mid_channels), nn.ReLU(inplaceTrue) ) self.conv_ca nn.Sequential( nn.Conv2d(in_channels, mid_channels, kernel_size3, padding1, biasFalse), nn.BatchNorm2d(mid_channels), nn.ReLU(inplaceTrue) ) # 双注意力模块 self.pa PositionAttention(mid_channels) self.ca ChannelAttention() # ChannelAttention内部无卷积输入输出通道一致 # 注意力后的融合与上采样卷积 self.conv_fusion nn.Sequential( nn.Conv2d(mid_channels, mid_channels, kernel_size3, padding1, biasFalse), nn.BatchNorm2d(mid_channels), nn.ReLU(inplaceTrue) ) # 最终分类卷积层 self.conv_cls nn.Sequential( nn.Conv2d(mid_channels, mid_channels // 2, kernel_size3, padding1, biasFalse), nn.BatchNorm2d(mid_channels // 2), nn.ReLU(inplaceTrue), nn.Dropout2d(0.1), # 可以加一个Dropout防止过拟合 nn.Conv2d(mid_channels // 2, num_classes, kernel_size1) # 1x1卷积将通道数映射为类别数 ) def forward(self, x): x: 主干网络输出的特征图形状 (B, in_channels, H//8, W//8) 输出: 分割预测图形状 (B, num_classes, H, W) # 分支处理 feat_pa self.conv_pa(x) feat_ca self.conv_ca(x) # 应用注意力 feat_pa_att self.pa(feat_pa) feat_ca_att self.ca(feat_ca) # 特征融合 feat_fusion feat_pa_att feat_ca_att feat_fusion self.conv_fusion(feat_fusion) # 分类预测 output self.conv_cls(feat_fusion) # 上采样回原图尺寸。输入是下采样8倍的特征需要上采样8倍。 # 通常使用双线性插值align_corners的设置需与数据预处理保持一致。 output F.interpolate(output, scale_factor8, modebilinear, align_cornersTrue) return output3. 搭建完整的DAnet结合主干网络与训练技巧有了强大的DAHead我们需要一个强壮的主干网络Backbone来提取基础特征。DAnet原文使用的是ResNet并对其进行了修改以适应密集预测任务。3.1 修改主干网络空洞卷积与特征提取标准的ResNet在最后两个阶段layer3和layer4会进行下采样导致最终输出的特征图尺寸是输入的1/32这对于需要精细边界的语义分割来说太粗糙了丢失了太多空间信息。常见的做法是引入空洞卷积Dilated Convolution。空洞卷积可以在不增加参数、不降低分辨率的情况下扩大卷积核的感受野。想象一下普通的3x3卷积它看的是相邻的9个点。如果设置空洞率dilation rate为2这个3x3的卷积核就会“跳过”一个像素去看实际感受野就变成了5x5。这就像你把一个密实的渔网拉松一些网眼变大覆盖的水面就更广但用的绳子参数还是那么多。在PyTorch中我们可以利用torchvision.models中的ResNet并通过replace_stride_with_dilation参数来轻松实现这个修改。import torchvision.models as models from torchvision.models._utils import IntermediateLayerGetter def get_backbone(backbone_nameresnet50, pretrainedTrue): 获取修改后的主干网络输出下采样8倍的特征图。 if backbone_name resnet50: backbone models.resnet50(pretrainedpretrained) # 关键修改将layer3和layer4的步幅替换为空洞卷积 # [False, True, True] 表示layer2的步幅不变layer3和layer4的步幅从2改为1并用空洞卷积扩大感受野 replace_stride_with_dilation [False, True, True] elif backbone_name resnet101: backbone models.resnet101(pretrainedpretrained) replace_stride_with_dilation [False, True, True] else: raise ValueError(fUnsupported backbone: {backbone_name}) # 移除最后的全连接层和平均池化层 backbone nn.Sequential(*list(backbone.children())[:-2]) # 更优雅的方式使用IntermediateLayerGetter直接获取指定层的输出 # 我们需要layer4的输出其下采样倍率为8如果替换了步幅 return_layers {layer4: out} backbone IntermediateLayerGetter(backbone, return_layersreturn_layers) return backbone # 测试一下 backbone get_backbone(resnet50, pretrainedFalse) dummy_input torch.randn(2, 3, 224, 224) with torch.no_grad(): features backbone(dummy_input) print(features[out].shape) # 期望输出torch.Size([2, 2048, 28, 28]) 224/828这样我们就得到了一个输出特征图为输入1/8大小的主干网络。这个特征图既保留了相对丰富的空间细节又拥有足够的高级语义信息非常适合作为双注意力头的输入。3.2 组装完整模型与训练配置将主干网络和DAHead组合起来就是完整的DAnet模型了。class DAnet(nn.Module): def __init__(self, num_classes, backboneresnet50, pretrainedTrue): super(DAnet, self).__init__() self.backbone get_backbone(backbone, pretrained) # 主干网络输出通道数ResNet50/101是2048 self.decoder DAHead(in_channels2048, num_classesnum_classes) def forward(self, x): # 主干网络提取特征 features self.backbone(x) # 返回一个字典 {out: tensor} # 解码器进行双注意力处理和上采样预测 out self.decoder(features[out]) return out # 实例化模型 model DAnet(num_classes21, backboneresnet50, pretrainedTrue) print(model)接下来是训练部分。语义分割是像素级分类所以损失函数通常选用交叉熵损失CrossEntropyLoss。对于像Cityscapes这类包含“忽略”类别如物体轮廓、遮挡边界的数据集可以在损失函数中设置ignore_index。优化器方面SGD with momentum 和 Adam 都是常见选择。论文中可能使用了SGD但根据我的经验使用Adam并配合合适的学习率衰减策略如CosineAnnealingLR或ReduceLROnPlateau通常更容易收敛且效果不错。一个非常重要的技巧是学习率热身Learning Rate Warmup。由于我们使用了在ImageNet上预训练的主干网络在训练初期直接使用较大的学习率可能会导致预训练权重被剧烈破坏模型不稳定。Warmup策略是在前几个epoch比如5个内让学习率从0线性增长到预设的初始值给模型一个“热身”的过程。import torch.optim as optim from torch.optim.lr_scheduler import CosineAnnealingLR, LinearLR from torch.nn import CrossEntropyLoss def get_optimizer_and_scheduler(model, lr1e-4, warmup_epochs5, total_epochs100): # 通常为骨干网络和分割头设置不同的学习率骨干网络的学习率小一些 backbone_params list(model.backbone.parameters()) decoder_params list(model.decoder.parameters()) optimizer optim.AdamW([ {params: backbone_params, lr: lr * 0.1}, # backbone学习率小10倍 {params: decoder_params, lr: lr} ], weight_decay1e-4) # 创建热身调度器 warmup_scheduler LinearLR(optimizer, start_factor0.01, end_factor1.0, total_iterswarmup_epochs) # 创建主调度器余弦退火 main_scheduler CosineAnnealingLR(optimizer, T_maxtotal_epochs - warmup_epochs, eta_minlr*1e-3) # 组合调度器先热身再余弦退火 from torch.optim.lr_scheduler import SequentialLR scheduler SequentialLR( optimizer, schedulers[warmup_scheduler, main_scheduler], milestones[warmup_epochs] ) criterion CrossEntropyLoss(ignore_index255) # 假设255为忽略的标签 return optimizer, scheduler, criterion数据增强对于语义分割模型至关重要它能极大地提升模型的泛化能力。推荐使用albumentations库它针对图像分割任务提供了丰富且高效的增强操作。import albumentations as A from albumentations.pytorch import ToTensorV2 def get_train_transform(): return A.Compose([ A.RandomResizedCrop(height512, width512, scale(0.5, 2.0)), # 随机裁剪和缩放 A.HorizontalFlip(p0.5), # 水平翻转 A.VerticalFlip(p0.1), # 垂直翻转 A.RandomRotate90(p0.5), # 随机旋转90度 A.OneOf([ # 随机选择一种颜色变换 A.HueSaturationValue(p1), A.RandomBrightnessContrast(p1), A.RandomGamma(p1), ], p0.5), A.Normalize(mean(0.485, 0.456, 0.406), std(0.229, 0.224, 0.225)), # ImageNet均值标准差 ToTensorV2(), # 转换为Tensor ]) def get_val_transform(): return A.Compose([ A.Resize(height512, width512), # 验证集只需resize到固定大小 A.Normalize(mean(0.485, 0.456, 0.406), std(0.229, 0.224, 0.225)), ToTensorV2(), ])4. 实战演练在CamVid数据集上训练与调试理论说得再多不如跑一遍代码来得实在。我们选用一个经典的语义分割基准数据集——CamVid驾驶场景来实战。这个数据集规模适中类别定义清晰非常适合做算法验证和教学。4.1 数据准备与加载首先你需要下载CamVid数据集并按照一定的目录结构放置。一个典型的目录结构如下camvid/ ├── train_images/ │ ├── 0001TP_006690.png │ └── ... ├── train_labels/ (对应的标注图像通常是单通道的索引图或RGB伪彩图) │ ├── 0001TP_006690.png │ └── ... ├── val_images/ ├── val_labels/ └── class_dict.csv (可选的类别颜色映射文件)然后我们编写一个PyTorch的Dataset类来加载数据。这里假设标签图像是RGB格式的伪彩色图我们需要将其转换为单通道的类别索引图。import os from PIL import Image import numpy as np import torch from torch.utils.data import Dataset, DataLoader class CamVidDataset(Dataset): def __init__(self, images_dir, masks_dir, transformNone, label_colorsNone): self.images_dir images_dir self.masks_dir masks_dir self.transform transform self.label_colors label_colors # 一个列表每个元素是一个(R,G,B)元组对应一个类别 self.images sorted(os.listdir(images_dir)) self.masks sorted(os.listdir(masks_dir)) assert len(self.images) len(self.masks), 图像和标签数量不匹配 def __len__(self): return len(self.images) def rgb_to_index(self, mask_rgb): 将RGB标签图转换为单通道类别索引图 index_mask np.zeros((mask_rgb.shape[0], mask_rgb.shape[1]), dtypenp.uint8) if self.label_colors is not None: for idx, color in enumerate(self.label_colors): # 找到所有像素值等于该颜色的位置赋值为类别索引 # 注意由于图像压缩颜色可能不完全相等需要设置一个容差 match np.all(np.abs(mask_rgb - color) 10, axis-1) index_mask[match] idx else: # 如果没提供颜色表简单将非零像素视为前景仅用于二分类演示 index_mask[(mask_rgb.sum(axis-1) 0)] 1 return index_mask def __getitem__(self, idx): img_path os.path.join(self.images_dir, self.images[idx]) mask_path os.path.join(self.masks_dir, self.masks[idx]) image np.array(Image.open(img_path).convert(RGB)) mask_rgb np.array(Image.open(mask_path).convert(RGB)) mask self.rgb_to_index(mask_rgb) # 转换为索引图 if self.transform: augmented self.transform(imageimage, maskmask) image augmented[image] mask augmented[mask].long() # 确保mask是Long类型 return image, mask # 定义CamVid的32个类别颜色示例需根据实际class_dict.csv填写 camvid_colors [ (64, 128, 64), # Animal (192, 0, 128), # Archway (0, 128, 192), # Bicyclist # ... 填写所有32个类别的RGB值 (0, 0, 0) # Void ] # 创建数据集和数据加载器 train_dataset CamVidDataset(camvid/train_images, camvid/train_labels, transformget_train_transform(), label_colorscamvid_colors) val_dataset CamVidDataset(camvid/val_images, camvid/val_labels, transformget_val_transform(), label_colorscamvid_colors) train_loader DataLoader(train_dataset, batch_size8, shuffleTrue, num_workers4, pin_memoryTrue) val_loader DataLoader(val_dataset, batch_size4, shuffleFalse, num_workers4, pin_memoryTrue)4.2 训练循环与模型评估训练循环是深度学习的标准流程但有几个细节需要注意混合精度训练使用torch.cuda.amp进行自动混合精度训练可以显著减少显存占用并加快训练速度尤其对于DAnet这种包含大量矩阵运算的模型收益明显。指标计算除了损失我们更关心语义分割的评估指标如平均交并比mIoU和像素准确率Pixel Accuracy。需要在每个epoch后对验证集进行计算。模型保存保存验证集上mIoU最高的模型而不是简单地按epoch保存。import torch.nn.functional as F from torch.cuda.amp import autocast, GradScaler from tqdm import tqdm def train_one_epoch(model, dataloader, optimizer, criterion, device, scalerNone, schedulerNone): model.train() total_loss 0.0 pbar tqdm(dataloader, descTraining) for images, masks in pbar: images, masks images.to(device), masks.to(device) optimizer.zero_grad() # 混合精度训练前向传播 with autocast(): outputs model(images) loss criterion(outputs, masks) # 混合精度训练反向传播 scaler.scale(loss).backward() scaler.step(optimizer) scaler.update() # 如果不用混合精度就是普通的 loss.backward() 和 optimizer.step() total_loss loss.item() pbar.set_postfix({loss: loss.item()}) if scheduler is not None: scheduler.step() return total_loss / len(dataloader) def evaluate(model, dataloader, criterion, device, num_classes): model.eval() total_loss 0.0 total_iou 0.0 total_pixel_acc 0.0 # 初始化混淆矩阵 conf_matrix torch.zeros(num_classes, num_classes).to(device) with torch.no_grad(): for images, masks in tqdm(dataloader, descEvaluating): images, masks images.to(device), masks.to(device) outputs model(images) loss criterion(outputs, masks) total_loss loss.item() # 获取预测类别 (B, C, H, W) - (B, H, W) preds outputs.argmax(dim1) # 计算像素准确率 correct (preds masks).sum().item() total_pixels masks.numel() - (masks 255).sum().item() # 忽略255标签 total_pixel_acc correct / total_pixels if total_pixels 0 else 0 # 更新混淆矩阵以计算mIoU (忽略255) for pred, mask in zip(preds.flatten(), masks.flatten()): if mask ! 255: # 忽略标签 conf_matrix[pred, mask] 1 # 计算mIoU # 交集对角线元素 intersection torch.diag(conf_matrix) # 并集预测和真实标签各自的总和减去交集 union conf_matrix.sum(dim0) conf_matrix.sum(dim1) - intersection iou_per_class intersection / (union 1e-8) # 避免除零 # 忽略背景类假设索引0是背景或计算所有类的平均 mean_iou iou_per_class[1:].mean().item() # 这里排除背景类 avg_loss total_loss / len(dataloader) avg_pixel_acc total_pixel_acc / len(dataloader) return avg_loss, avg_pixel_acc, mean_iou # 主训练循环 device torch.device(cuda if torch.cuda.is_available() else cpu) model DAnet(num_classeslen(camvid_colors)).to(device) optimizer, scheduler, criterion get_optimizer_and_scheduler(model) scaler GradScaler() # 用于混合精度训练 best_miou 0.0 num_epochs 100 for epoch in range(num_epochs): print(fEpoch {epoch1}/{num_epochs}) train_loss train_one_epoch(model, train_loader, optimizer, criterion, device, scaler, scheduler) val_loss, val_pixel_acc, val_miou evaluate(model, val_loader, criterion, device, len(camvid_colors)) print(fTrain Loss: {train_loss:.4f} | Val Loss: {val_loss:.4f} | Val Pixel Acc: {val_pixel_acc:.4f} | Val mIoU: {val_miou:.4f}) # 保存最佳模型 if val_miou best_miou: best_miou val_miou torch.save({ epoch: epoch, model_state_dict: model.state_dict(), optimizer_state_dict: optimizer.state_dict(), val_miou: val_miou, }, best_danet_camvid.pth) print(fBest model saved with mIoU: {val_miou:.4f})4.3 可视化与问题排查训练过程中可视化是理解模型在做什么、哪里出问题的关键手段。我们可以可视化以下几样东西注意力图在DAHead的forward方法中可以额外返回空间注意力图attention_s和通道注意力图attention_x。将它们reshape并可视化可以看到模型关注了哪些空间位置和特征通道。例如空间注意力图经过reshape后是(N, N)我们可以取其中一行对应一个像素点将其reshape回(H, W)就能看到所有像素对这个点的“贡献”热力图。预测结果定期将验证集的原始图像、真实标签和模型预测结果并排显示。这能直观地看到模型分割的边界是否清晰小物体是否被漏掉类别混淆发生在哪里。损失曲线与指标曲线绘制训练损失、验证损失、mIoU随epoch的变化曲线。如果训练损失持续下降但验证损失上升很可能过拟合了。如果mIoU很早就停滞不前可能需要检查学习率、数据增强或模型容量。在实现时我踩过的一个坑是关于矩阵乘法的维度和转置。在实现空间注意力时B^T * C与D * S^T这两个乘法的顺序和转置关系一定要对照论文公式和图示反复确认否则得到的注意力机制效果会大打折扣甚至起反作用。另一个常见的坑是上采样时的align_corners参数。在PyTorch的F.interpolate函数中align_corners决定了像素网格的对齐方式。如果预处理时使用的库如OpenCV, PIL, albumentations和模型上采样时设置不一致会导致特征图出现几个像素的错位严重影响边界精度。我的经验是如果预处理用了albumentations.Resize上采样时设置align_cornersFalse通常能保持一致。最后别忘了利用TensorBoard或WandB等工具来记录所有的训练指标、图像和注意力图这对于分析和调试模型有巨大的帮助。训练一个像DAnet这样的模型在CamVid上达到不错的mIoU比如70%以上通常需要几十到上百个epoch具体取决于你的硬件和超参数设置。耐心调整观察曲线分析bad case你就能逐渐驾驭这个强大的双注意力分割模型。