拓冰网站定制 
红外小目标检测这个方向做的人不少但真正能把公开网络复现到论文指标能打平、小目标不丢、工程能跑通的程度中间踩的坑远比想象中多。DNANetDense Nested Attention Network是红外小目标检测里比较有代表性的一套结构核心卖点就是通过密集嵌套连接加注意力机制把深层语义和浅层细节反复融合缓解小目标在多次下采样后消失的问题。我最近用PyTorch把这套网络从零复现了一遍从环境搭建、数据准备、模块实现到训练调参、推理验证完整走了一遍流程中间也踩了不少坑。这篇就把整个复现过程拆开讲清楚包括每一步为什么这么做、参数怎么定、哪些地方最容易翻车适合有一定PyTorch基础、想上手红外小目标检测或者想复现这篇网络的朋友参考。整套代码基于PyTorch实现环境用Anaconda管理训练和推理都在单卡上完成下面按实际操作的顺序展开。1. 先搞清楚DNANet到底在解决什么问题1.1 红外小目标的小到底有多小红外小目标检测和常规目标检测最大的区别在于目标本身的尺度。常规检测里一个目标可能占几百上千像素而红外小目标往往只有几个到几十个像素在整张图里占比可能不到千分之一。这种尺度下目标没有明显的纹理、形状、颜色特征很多时候就是一团比背景稍亮的模糊斑点信噪比极低。这就带来一个很现实的问题常规检测网络靠多次下采样扩大感受野但每下采样一次小目标的特征就被稀释一次。等到深层特征图目标可能只剩一两个像素甚至直接消失。这就是标题里说的小目标丢失——不是网络不work而是目标在特征传递过程中被洗掉了。DNANet的设计动机就来自这里。它不去追求更深的骨干网络而是把重点放在如何让浅层的高分辨率细节一路保留到深层并且和深层语义反复交互上。理解了这一点后面所有模块的设计就都顺了。1.2 密集嵌套连接和注意力各自的作用DNANet的核心可以拆成两块密集嵌套结构Dense Nested和注意力模块。密集嵌套解决的是特征复用和梯度流动问题注意力解决的是在杂乱背景里把真正的小目标拎出来的问题。密集嵌套的思路是每一层的输出不只传给下一层还传给后面所有层同时每一层也接收前面所有层的输入。这样浅层的高分辨率特征可以跳过中间层直接参与深层计算小目标的细节不会被逐层磨掉。这和DenseNet的思路一脉相承但DNANet把它用在了编解码的嵌套结构里融合的粒度更细。注意力模块则负责在融合后的特征上做加权。红外图像背景复杂云层边缘、地面热源、建筑反光都可能产生类似小目标的亮斑注意力机制通过通道和空间两个维度重新分配权重把真正目标的响应放大把背景杂波压下去。这两块配合起来才是DNANet能不丢小目标的根本原因。1.3 复现前必须明确的几个预期在动手之前有几个预期要先摆正否则很容易中途怀疑人生。第一DNANet不是即插即用的通用检测器它是针对红外小目标这个特定场景设计的换到常规目标检测数据集上未必有优势。第二复现的目标是结构对齐、指标接近不是逐比特一致因为论文里很多训练细节学习率调度、数据增强的具体参数不会写全需要自己补。第三小目标检测对数据质量极其敏感如果数据集本身标注质量差再好的网络也救不回来。把这三条想清楚后面的工作就是按部就班地搭环境、写模块、调参、验证。2. 环境搭建Anaconda PyTorch 的稳妥配置2.1 为什么用conda而不是直接pip很多人装PyTorch习惯直接pip install简单是简单但红外检测这类项目往往要装一堆科学计算和图像处理库版本冲突的概率很高。conda的优势在于它能管理非Python依赖比如CUDA相关的底层库环境隔离也更干净。一旦某个库版本不对直接删环境重建不会污染系统Python。我的做法是给这个项目单独建一个环境命名上带项目标识方便以后同时维护多个复现项目。具体命令如下conda create -n dnanet python3.9 -y conda activate dnanetPython版本选3.9是比较稳的选择3.10以上有些老库的wheel还没跟上3.8又偏旧3.9在兼容性和新特性之间平衡得最好。2.2 PyTorch版本和CUDA的匹配这一步是环境搭建里最容易翻车的地方。PyTorch版本、CUDA版本、显卡驱动版本三者必须匹配否则要么装完import报错要么能import但一跑就提示CUDA不可用。我的建议是先去PyTorch官网的安装页面根据自己的CUDA驱动版本选对应的安装命令。查看驱动支持的最高CUDA版本可以用nvidia-smi右上角会显示CUDA Version: xx.x这个是你驱动支持的最高版本装PyTorch时选的CUDA版本不能超过它。比如显示12.1那你可以装cu118或cu121的PyTorch但不能装cu124的。假设选CUDA 11.8安装命令大致是conda install pytorch torchvision torchaudio pytorch-cuda11.8 -c pytorch -c nvidia装完之后一定要验证别急着往下走import torch print(torch.__version__) print(torch.cuda.is_available()) print(torch.cuda.get_device_name(0))三个输出分别是版本号、True、显卡型号才算真正装好。如果is_available()返回False八成是版本不匹配回退重装。提示如果手头只有CPU也能跑通整个流程只是训练会慢很多。把上面命令里的pytorch-cuda去掉装CPU版即可代码里device判断写成自动选择就行。2.3 其余依赖库的安装除了PyTorch这个项目还需要几个常规库numpy、opencv-python用于图像读写、tqdm用于进度显示、matplotlib用于可视化、scikit-image用于一些图像指标计算。这些直接用pip装就行conda装图像库有时候反而慢。pip install numpy opencv-python tqdm matplotlib scikit-image装完之后建议跑一个最小验证脚本把torch、cv2、numpy都import一遍确认没有报错。这一步花两分钟能省掉后面调试时排查环境问题的大量时间。2.4 目录结构规划环境好了之后先把项目目录规划清楚。我习惯按功能分目录避免所有文件堆在一起dnanet/ ├── data/ # 数据集 │ ├── train/ │ │ ├── images/ │ │ └── masks/ │ └── test/ │ ├── images/ │ └── masks/ ├── models/ # 网络结构 │ ├── dnanet.py │ └── modules.py ├── utils/ # 工具函数 │ ├── dataset.py │ ├── metrics.py │ └── visualize.py ├── train.py ├── test.py └── config.py这种结构的好处是职责清晰改网络只动models改数据加载只动utils训练脚本保持干净。config.py单独放超参数调参时不用在代码里到处找。3. 数据准备红外小目标数据集的坑3.1 数据集的选择与格式红外小目标检测常用的公开数据集有几个比如NUAA-SIRST、NUDT-SIRST、IRSTD-1k等。这些数据集基本都是图像二值掩码的形式掩码里目标区域为1背景为0。选数据集时要注意两点一是标注质量有些数据集的目标标注偏大或偏小会直接影响训练二是场景多样性如果全是天空背景模型换到地面场景就废了。我这次用的是NUAA-SIRST图像是单通道灰度图尺寸不统一掩码是同样尺寸的二值图。数据量不大几百张所以数据增强必须做足否则很容易过拟合。3.2 图像和掩码的配对读取数据加载最容易出的问题是图像和掩码对不上。我的做法是让图像和掩码文件名完全一致只是放在不同目录加载时按文件名匹配。dataset.py的核心逻辑大致是这样import os import cv2 import torch import numpy as np from torch.utils.data import Dataset class IRSTDDataset(Dataset): def __init__(self, img_dir, mask_dir, transformNone): self.img_dir img_dir self.mask_dir mask_dir self.transform transform self.names sorted(os.listdir(img_dir)) def __len__(self): return len(self.names) def __getitem__(self, idx): name self.names[idx] img_path os.path.join(self.img_dir, name) mask_path os.path.join(self.mask_dir, name) img cv2.imread(img_path, cv2.IMREAD_GRAYSCALE) mask cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) img img.astype(np.float32) / 255.0 mask (mask 127).astype(np.float32) img np.expand_dims(img, axis0) mask np.expand_dims(mask, axis0) if self.transform: img, mask self.transform(img, mask) return torch.from_numpy(img), torch.from_numpy(mask)这里有几个细节值得说。灰度图用IMREAD_GRAYSCALE读保证是单通道归一化到0-1避免数值范围差异影响训练掩码用阈值127二值化因为有些掩码保存时边缘会有灰度过渡最后都加一个通道维度变成(C, H, W)格式符合PyTorch的输入约定。3.3 数据增强小目标场景的特殊处理常规数据增强里的随机裁剪、缩放对小目标其实是有风险的。你随机裁一块很可能把本来就没几个像素的目标裁掉或者缩放后目标变得更小。所以小目标检测的数据增强要更克制。我用的增强组合是随机水平翻转、随机垂直翻转、随机旋转90度的整数倍。这三种操作不会改变目标的尺度也不会把目标裁掉相对安全。如果要加亮度扰动幅度也要小因为红外图像的灰度本身就携带了目标的强度信息扰动太大会破坏这个特征。import random def augment(img, mask): if random.random() 0.5: img np.flip(img, axis2).copy() mask np.flip(mask, axis2).copy() if random.random() 0.5: img np.flip(img, axis1).copy() mask np.flip(mask, axis1).copy() k random.choice([0, 1, 2, 3]) if k 0: img np.rot90(img, k, axes(1, 2)).copy() mask np.rot90(mask, k, axes(1, 2)).copy() return img, mask注意flip和rot90之后要.copy()否则numpy返回的是视图后续操作可能出问题。这个坑我踩过一次表现为训练时loss莫名其妙不下降排查半天才发现是数据增强返回了共享内存的视图。3.4 训练集和测试集的划分数据量小的时候划分要讲究。我按8:2划分训练和测试并且保证两个集合里都包含各种场景天空、地面、海面。如果随机划分导致测试集全是简单场景评估结果会虚高没有参考价值。划分完之后建议把训练集里目标像素占比统计一下。如果某些图里目标像素占比极低比如低于0.01%训练时这些图的loss贡献几乎为零可以考虑适当过采样或者用带权重的损失函数补偿。4. DNANet网络结构的逐模块实现4.1 整体架构的拆解思路DNANet整体是编码器-解码器结构编码器逐层下采样提取特征解码器逐层上采样恢复分辨率中间用密集嵌套连接把各层特征反复融合。和U-Net的区别在于U-Net的跳跃连接是一对一的而DNANet是多对多的每一层都接收前面所有层的特征。实现的时候我建议先把整体骨架搭出来用占位模块跑通前向传播确认张量维度都对再逐个替换成真正的模块。这样调试成本最低不会一上来就被一堆维度错误淹没。4.2 基础卷积块的设计基础卷积块是网络的积木设计上要兼顾表达能力和计算量。我用的是卷积BN激活的组合激活函数选LeakyReLU而不是ReLU原因是红外小目标的响应本来就弱ReLU在负半轴直接置零容易把弱信号彻底掐掉LeakyReLU保留一点负半轴斜率对小目标更友好。import torch.nn as nn class ConvBlock(nn.Module): def __init__(self, in_ch, out_ch, kernel_size3, stride1): super().__init__() padding kernel_size // 2 self.conv nn.Conv2d(in_ch, out_ch, kernel_size, stride, padding, biasFalse) self.bn nn.BatchNorm2d(out_ch) self.act nn.LeakyReLU(0.1, inplaceTrue) def forward(self, x): return self.act(self.bn(self.conv(x)))bias设为False是因为后面接了BNBN本身有偏置再叠加卷积的bias是冗余的还多占显存。这个细节很多人不注意但在小显存卡上积少成多。4.3 密集嵌套融合模块的实现这是DNANet的核心。密集嵌套的意思是解码器第i层的输入来自编码器对应层和解码器前面所有层输出的拼接。实现上用一个列表缓存各层输出每次融合时把需要的层拼起来。class DenseFusion(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.conv ConvBlock(in_ch, out_ch) def forward(self, feats): x torch.cat(feats, dim1) return self.conv(x)调用的时候把当前层需要的所有特征放进feats列表cat之后过一层卷积降维。这里要注意通道数的计算拼接后的通道数是所有输入通道之和ConvBlock的in_ch必须和这个和一致否则直接报维度错误。我一般会在config里把每层的通道数写清楚避免手算出错。密集嵌套带来的一个副作用是显存占用上升因为要缓存多层特征。如果显存吃紧可以适当减少融合的层数或者用1x1卷积先降维再拼接。4.4 注意力模块CSAM与DNIM的配合注意力部分是DNANet区别于普通密集网络的关键。这里涉及两个模块一个是通道-空间注意力CSAM一个是密集嵌套交互模块DNIM两者配合完成特征的重新加权。CSAM的思路是先做通道注意力再做空间注意力串行叠加。通道注意力回答哪些通道更重要空间注意力回答图上哪些位置更重要。实现上通道注意力用全局平均池化加两层全连接空间注意力用通道维度的最大池化和平均池化拼接后过一层卷积。class CSAM(nn.Module): def __init__(self, channels, reduction16): super().__init__() self.avg_pool nn.AdaptiveAvgPool2d(1) self.fc nn.Sequential( nn.Linear(channels, channels // reduction, biasFalse), nn.ReLU(inplaceTrue), nn.Linear(channels // reduction, channels, biasFalse), nn.Sigmoid() ) self.spatial nn.Conv2d(2, 1, kernel_size7, padding3, biasFalse) self.sigmoid nn.Sigmoid() def forward(self, x): b, c, _, _ x.size() y self.avg_pool(x).view(b, c) y self.fc(y).view(b, c, 1, 1) x x * y avg_out torch.mean(x, dim1, keepdimTrue) max_out, _ torch.max(x, dim1, keepdimTrue) s torch.cat([avg_out, max_out], dim1) s self.sigmoid(self.spatial(s)) return x * sreduction取16是常规做法通道数少的时候可以取8避免降维后信息损失太多。空间注意力的卷积核用7x7感受野够大能覆盖小目标周围的上下文。DNIM则负责在密集嵌套的路径上做交互把不同层级的特征通过注意力加权后再融合。它的实现比CSAM复杂一些核心是把相邻层的特征做交叉注意力让深层语义指导浅层特征的筛选同时浅层细节补充深层的空间信息。注意注意力模块不要堆太多。我一开始在每一层融合后都加了CSAM结果显存爆了不说训练还变得很不稳定。后来改成只在关键的三层融合后加效果反而更好。注意力是双刃剑加多了会过拟合尤其是在小数据集上。4.5 上采样与输出层解码器的上采样用双线性插值而不是转置卷积。转置卷积容易产生棋盘格伪影对小目标检测是致命的因为伪影本身就像小亮点会干扰判断。双线性插值平滑没有这个问题。class UpBlock(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.up nn.Upsample(scale_factor2, modebilinear, align_cornersTrue) self.conv ConvBlock(in_ch, out_ch) def forward(self, x): return self.conv(self.up(x))输出层用1x1卷积把通道数降到1然后过Sigmoid得到0-1的概率图。训练时和掩码算二值交叉熵推理时按阈值二值化。5. 训练流程与损失函数的选择5.1 为什么不能只用BCE红外小目标检测有个天然的类别不平衡问题目标像素可能只占全图的千分之一甚至更少背景像素占绝大多数。如果只用二值交叉熵BCE网络会发现全部预测为背景就能拿到很低的loss于是学不到任何东西这就是所谓的背景主导。解决办法是组合损失。我用的方案是BCE加Dice损失BCE负责逐像素的稳定梯度Dice负责优化整体的重叠度两者互补。Dice对小目标尤其友好因为它关注的是预测和真值的重叠比例而不是绝对像素数。class BCEDiceLoss(nn.Module): def __init__(self, bce_weight0.5): super().__init__() self.bce_weight bce_weight self.bce nn.BCELoss() def forward(self, pred, target): bce_loss self.bce(pred, target) smooth 1e-6 pred_flat pred.view(-1) target_flat target.view(-1) intersection (pred_flat * target_flat).sum() dice_loss 1 - (2 * intersection smooth) / \ (pred_flat.sum() target_flat.sum() smooth) return self.bce_weight * bce_loss \ (1 - self.bce_weight) * dice_lossbce_weight取0.5是个稳妥的起点如果发现召回率偏低漏检多可以调低bce权重让Dice占主导如果误检多就调高bce权重。5.2 优化器和学习率的设置优化器用Adam初始学习率1e-3。Adam对学习率不那么敏感适合快速起步。但训练到后期固定学习率会导致loss在最优值附近震荡所以要用学习率调度。我用的是余弦退火加warmup。warmup在前几个epoch把学习率从很小的值线性升到初始值避免一开始梯度太大把网络带偏余弦退火让学习率平滑下降到接近零帮助收敛到更优的点。from torch.optim.lr_scheduler import CosineAnnealingLR, LambdaLR def get_scheduler(optimizer, warmup_epochs, total_epochs): def lr_lambda(epoch): if epoch warmup_epochs: return (epoch 1) / warmup_epochs progress (epoch - warmup_epochs) / \ (total_epochs - warmup_epochs) return 0.5 * (1 np.cos(np.pi * progress)) return LambdaLR(optimizer, lr_lambda)warmup一般设3-5个epoch总epoch数根据数据量定小数据集100-200个epoch比较合适。5.3 训练循环的关键细节训练循环本身不复杂但有几个细节决定成败。第一每个epoch开始前要调用model.train()结束评估时调用model.eval()并且用torch.no_grad()包住评估否则BN和Dropout的行为不对显存也会白白占用。第二梯度要清零。PyTorch默认累积梯度忘记zero_grad()会导致梯度越积越大训练直接崩。第三如果显存不够用梯度累积模拟大batch。比如想要batch size 16但显存只够4就累积4次梯度再更新一次。for epoch in range(total_epochs): model.train() for img, mask in train_loader: img, mask img.to(device), mask.to(device) pred model(img) loss criterion(pred, mask) optimizer.zero_grad() loss.backward() optimizer.step() scheduler.step() # 评估 model.eval() with torch.no_grad(): for img, mask in test_loader: img, mask img.to(device), mask.to(device) pred model(img) # 计算指标5.4 训练过程中的监控指标光看loss不够loss下降不代表检测效果好。要同时监控几个指标IoU交并比、Pd检测率即召回率、Fa虚警率。IoU衡量整体重叠度Pd衡量漏检情况Fa衡量误检情况。小目标检测里Pd和Fa往往是一对矛盾调参就是在两者之间找平衡。我习惯每个epoch把这三个指标打印出来画成曲线。如果Pd上不去说明网络没学到目标如果Fa太高说明背景杂波没压住可能要加注意力或者调损失权重。6. 推理与后处理让结果真正可用6.1 推理流程和阈值选择推理时把测试图送进网络得到0-1的概率图然后按阈值二值化。阈值的选择很关键太高会漏检太低会误检。默认0.5不一定最优我一般会在验证集上扫一遍阈值从0.1到0.9看哪个阈值下F1最高。def inference(model, img, threshold0.5): model.eval() with torch.no_grad(): img img.to(device) pred model(img) pred (pred threshold).float() return pred如果测试集和训练集分布差异大阈值可能要重新调。这个没有一劳永逸的办法只能根据实际数据来。6.2 连通域后处理去噪网络输出的二值图里往往有一些孤立的噪点这些噪点面积很小明显不是目标。可以用连通域分析过滤掉面积过小的区域。import cv2 import numpy as np def remove_small_regions(mask, min_area3): mask mask.astype(np.uint8) num, labels, stats, _ cv2.connectedComponentsWithStats(mask, 8) out np.zeros_like(mask) for i in range(1, num): if stats[i, cv2.CC_STAT_AREA] min_area: out[labels i] 1 return outmin_area要根据目标尺度定如果目标最小只有几个像素min_area设太大反而会把真目标滤掉。我一般设3-5具体看数据集。6.3 可视化验证推理完一定要可视化把原图、真值掩码、预测掩码并排画出来。光看指标数字容易自欺欺人可视化能一眼看出问题是漏检了、误检了还是边界对不齐。import matplotlib.pyplot as plt def visualize(img, gt, pred): fig, axes plt.subplots(1, 3, figsize(12, 4)) axes[0].imshow(img[0], cmapgray) axes[0].set_title(Image) axes[1].imshow(gt[0], cmapgray) axes[1].set_title(Ground Truth) axes[2].imshow(pred[0], cmapgray) axes[2].set_title(Prediction) for ax in axes: ax.axis(off) plt.show()我习惯每次评估都随机抽几张图看尤其是那些指标差的图往往能发现系统性问题。7. 复现过程中踩过的坑与调参经验7.1 小目标在深层特征图消失这是复现初期最典型的问题。训练了几十个epochloss降下去了但预测图里小目标几乎全丢。排查下来发现是下采样太激进目标在深层特征图上已经不足一个像素。解决办法有两个一是减少下采样次数把编码器从5层减到4层二是在密集嵌套里强化浅层特征的权重让高分辨率特征更多地参与深层计算。我两个都用了效果明显改善。7.2 损失不下降的几种原因loss不下降的原因很多我按排查顺序列一下先看数据加载对不对把batch里的图和掩码可视化出来确认配对没错再看学习率是不是太大太大导致震荡太小导致不动然后看损失函数实现有没有bug尤其是Dice的分母有没有加smooth防止除零最后看网络输出有没有经过Sigmoid如果输出是logits却直接喂给BCELossloss会异常。7.3 过拟合的识别与缓解小数据集上过拟合几乎是必然的。识别方法很简单训练loss持续下降但验证指标到某个点后开始变差。缓解手段按优先级排先加数据增强再加Dropout或权重衰减最后考虑减小模型容量。我在全连接层前加了Dropout(0.3)并在优化器里设了weight_decay1e-4过拟合明显缓解。但要注意Dropout不要加在卷积层后面会破坏空间信息。7.4 显存不足的应对策略显存不足时按这个顺序优化先减小batch size这是最直接的再检查有没有不必要的中间变量占用显存比如在训练循环里保留了计算图然后考虑混合精度训练用torch.cuda.amp能省将近一半显存最后才是减小模型。混合精度训练的实现from torch.cuda.amp import autocast, GradScaler scaler GradScaler() for img, mask in train_loader: optimizer.zero_grad() with autocast(): pred model(img) loss criterion(pred, mask) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()混合精度要注意某些操作在fp16下会溢出如果发现loss变成nan先关掉amp排查。7.5 指标波动的正常范围训练过程中指标有波动是正常的尤其是小数据集。如果Pd在几个百分点内波动不用太在意如果波动超过10个百分点说明训练不稳定要检查学习率、batch size或者数据增强的随机性。我一般会跑多次实验取平均单次结果参考价值有限。如果时间紧至少固定随机种子跑一次保证结果可复现。8. 从复现到改进几个可以尝试的方向8.1 骨干网络的替换DNANet原版用的骨干比较轻量如果想提升性能可以换成更强的骨干比如ResNet系列。但要注意骨干越深小目标消失的风险越大所以换骨干的同时要强化密集嵌套的融合或者引入多尺度特征。8.2 损失函数的进一步优化BCE加Dice是基础组合还可以尝试Focal Loss或者Tversky Loss。Focal Loss通过调制因子让网络更关注难样本对小目标有帮助Tversky Loss可以调节假阳和假阴的权重适合对漏检或误检有偏好的场景。8.3 多尺度训练与测试小目标的尺度变化是个难题。多尺度训练指在训练时随机缩放输入让网络见到不同尺度的目标多尺度测试指在推理时对同一张图用多个尺度预测再融合。这两个技巧都能提升尺度鲁棒性代价是训练和推理时间增加。8.4 部署时的优化考虑如果最终要部署推理速度是个硬指标。可以做的优化包括把BN层折叠进卷积、用TensorRT或ONNX Runtime加速、量化到INT8。这些优化对精度有轻微影响需要重新评估指标。我在实际部署时发现BN折叠几乎无损量化到INT8会掉一两个点的IoU如果对精度要求高建议只做BN折叠和算子融合。整套流程走下来DNANet的复现难点其实不在网络结构本身而在数据质量、损失设计和训练细节的把控上。网络结构照着论文搭维度对齐了就能跑但要让小目标真正不丢得在数据增强、损失权重、注意力模块的取舍上反复试。我个人的体会是先把baseline跑通指标能到论文的八九成再逐步加改进比一上来就堆模块要稳得多。另外红外小目标这个方向对数据的依赖极强如果手头数据场景单一建议先扩充数据或者做场景迁移否则调参调到天荒地老也难有突破。 PyTorch从零复现DNANet：红外小目标检测实战与调参避坑指南
红外小目标检测这个方向做的人不少但真正能把公开网络复现到论文指标能打平、小目标不丢、工程能跑通的程度中间踩的坑远比想象中多。DNANetDense Nested Attention Network是红外小目标检测里比较有代表性的一套结构核心卖点就是通过密集嵌套连接加注意力机制把深层语义和浅层细节反复融合缓解小目标在多次下采样后消失的问题。我最近用PyTorch把这套网络从零复现了一遍从环境搭建、数据准备、模块实现到训练调参、推理验证完整走了一遍流程中间也踩了不少坑。这篇就把整个复现过程拆开讲清楚包括每一步为什么这么做、参数怎么定、哪些地方最容易翻车适合有一定PyTorch基础、想上手红外小目标检测或者想复现这篇网络的朋友参考。整套代码基于PyTorch实现环境用Anaconda管理训练和推理都在单卡上完成下面按实际操作的顺序展开。1. 先搞清楚DNANet到底在解决什么问题1.1 红外小目标的小到底有多小红外小目标检测和常规目标检测最大的区别在于目标本身的尺度。常规检测里一个目标可能占几百上千像素而红外小目标往往只有几个到几十个像素在整张图里占比可能不到千分之一。这种尺度下目标没有明显的纹理、形状、颜色特征很多时候就是一团比背景稍亮的模糊斑点信噪比极低。这就带来一个很现实的问题常规检测网络靠多次下采样扩大感受野但每下采样一次小目标的特征就被稀释一次。等到深层特征图目标可能只剩一两个像素甚至直接消失。这就是标题里说的小目标丢失——不是网络不work而是目标在特征传递过程中被洗掉了。DNANet的设计动机就来自这里。它不去追求更深的骨干网络而是把重点放在如何让浅层的高分辨率细节一路保留到深层并且和深层语义反复交互上。理解了这一点后面所有模块的设计就都顺了。1.2 密集嵌套连接和注意力各自的作用DNANet的核心可以拆成两块密集嵌套结构Dense Nested和注意力模块。密集嵌套解决的是特征复用和梯度流动问题注意力解决的是在杂乱背景里把真正的小目标拎出来的问题。密集嵌套的思路是每一层的输出不只传给下一层还传给后面所有层同时每一层也接收前面所有层的输入。这样浅层的高分辨率特征可以跳过中间层直接参与深层计算小目标的细节不会被逐层磨掉。这和DenseNet的思路一脉相承但DNANet把它用在了编解码的嵌套结构里融合的粒度更细。注意力模块则负责在融合后的特征上做加权。红外图像背景复杂云层边缘、地面热源、建筑反光都可能产生类似小目标的亮斑注意力机制通过通道和空间两个维度重新分配权重把真正目标的响应放大把背景杂波压下去。这两块配合起来才是DNANet能不丢小目标的根本原因。1.3 复现前必须明确的几个预期在动手之前有几个预期要先摆正否则很容易中途怀疑人生。第一DNANet不是即插即用的通用检测器它是针对红外小目标这个特定场景设计的换到常规目标检测数据集上未必有优势。第二复现的目标是结构对齐、指标接近不是逐比特一致因为论文里很多训练细节学习率调度、数据增强的具体参数不会写全需要自己补。第三小目标检测对数据质量极其敏感如果数据集本身标注质量差再好的网络也救不回来。把这三条想清楚后面的工作就是按部就班地搭环境、写模块、调参、验证。2. 环境搭建Anaconda PyTorch 的稳妥配置2.1 为什么用conda而不是直接pip很多人装PyTorch习惯直接pip install简单是简单但红外检测这类项目往往要装一堆科学计算和图像处理库版本冲突的概率很高。conda的优势在于它能管理非Python依赖比如CUDA相关的底层库环境隔离也更干净。一旦某个库版本不对直接删环境重建不会污染系统Python。我的做法是给这个项目单独建一个环境命名上带项目标识方便以后同时维护多个复现项目。具体命令如下conda create -n dnanet python3.9 -y conda activate dnanetPython版本选3.9是比较稳的选择3.10以上有些老库的wheel还没跟上3.8又偏旧3.9在兼容性和新特性之间平衡得最好。2.2 PyTorch版本和CUDA的匹配这一步是环境搭建里最容易翻车的地方。PyTorch版本、CUDA版本、显卡驱动版本三者必须匹配否则要么装完import报错要么能import但一跑就提示CUDA不可用。我的建议是先去PyTorch官网的安装页面根据自己的CUDA驱动版本选对应的安装命令。查看驱动支持的最高CUDA版本可以用nvidia-smi右上角会显示CUDA Version: xx.x这个是你驱动支持的最高版本装PyTorch时选的CUDA版本不能超过它。比如显示12.1那你可以装cu118或cu121的PyTorch但不能装cu124的。假设选CUDA 11.8安装命令大致是conda install pytorch torchvision torchaudio pytorch-cuda11.8 -c pytorch -c nvidia装完之后一定要验证别急着往下走import torch print(torch.__version__) print(torch.cuda.is_available()) print(torch.cuda.get_device_name(0))三个输出分别是版本号、True、显卡型号才算真正装好。如果is_available()返回False八成是版本不匹配回退重装。提示如果手头只有CPU也能跑通整个流程只是训练会慢很多。把上面命令里的pytorch-cuda去掉装CPU版即可代码里device判断写成自动选择就行。2.3 其余依赖库的安装除了PyTorch这个项目还需要几个常规库numpy、opencv-python用于图像读写、tqdm用于进度显示、matplotlib用于可视化、scikit-image用于一些图像指标计算。这些直接用pip装就行conda装图像库有时候反而慢。pip install numpy opencv-python tqdm matplotlib scikit-image装完之后建议跑一个最小验证脚本把torch、cv2、numpy都import一遍确认没有报错。这一步花两分钟能省掉后面调试时排查环境问题的大量时间。2.4 目录结构规划环境好了之后先把项目目录规划清楚。我习惯按功能分目录避免所有文件堆在一起dnanet/ ├── data/ # 数据集 │ ├── train/ │ │ ├── images/ │ │ └── masks/ │ └── test/ │ ├── images/ │ └── masks/ ├── models/ # 网络结构 │ ├── dnanet.py │ └── modules.py ├── utils/ # 工具函数 │ ├── dataset.py │ ├── metrics.py │ └── visualize.py ├── train.py ├── test.py └── config.py这种结构的好处是职责清晰改网络只动models改数据加载只动utils训练脚本保持干净。config.py单独放超参数调参时不用在代码里到处找。3. 数据准备红外小目标数据集的坑3.1 数据集的选择与格式红外小目标检测常用的公开数据集有几个比如NUAA-SIRST、NUDT-SIRST、IRSTD-1k等。这些数据集基本都是图像二值掩码的形式掩码里目标区域为1背景为0。选数据集时要注意两点一是标注质量有些数据集的目标标注偏大或偏小会直接影响训练二是场景多样性如果全是天空背景模型换到地面场景就废了。我这次用的是NUAA-SIRST图像是单通道灰度图尺寸不统一掩码是同样尺寸的二值图。数据量不大几百张所以数据增强必须做足否则很容易过拟合。3.2 图像和掩码的配对读取数据加载最容易出的问题是图像和掩码对不上。我的做法是让图像和掩码文件名完全一致只是放在不同目录加载时按文件名匹配。dataset.py的核心逻辑大致是这样import os import cv2 import torch import numpy as np from torch.utils.data import Dataset class IRSTDDataset(Dataset): def __init__(self, img_dir, mask_dir, transformNone): self.img_dir img_dir self.mask_dir mask_dir self.transform transform self.names sorted(os.listdir(img_dir)) def __len__(self): return len(self.names) def __getitem__(self, idx): name self.names[idx] img_path os.path.join(self.img_dir, name) mask_path os.path.join(self.mask_dir, name) img cv2.imread(img_path, cv2.IMREAD_GRAYSCALE) mask cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) img img.astype(np.float32) / 255.0 mask (mask 127).astype(np.float32) img np.expand_dims(img, axis0) mask np.expand_dims(mask, axis0) if self.transform: img, mask self.transform(img, mask) return torch.from_numpy(img), torch.from_numpy(mask)这里有几个细节值得说。灰度图用IMREAD_GRAYSCALE读保证是单通道归一化到0-1避免数值范围差异影响训练掩码用阈值127二值化因为有些掩码保存时边缘会有灰度过渡最后都加一个通道维度变成(C, H, W)格式符合PyTorch的输入约定。3.3 数据增强小目标场景的特殊处理常规数据增强里的随机裁剪、缩放对小目标其实是有风险的。你随机裁一块很可能把本来就没几个像素的目标裁掉或者缩放后目标变得更小。所以小目标检测的数据增强要更克制。我用的增强组合是随机水平翻转、随机垂直翻转、随机旋转90度的整数倍。这三种操作不会改变目标的尺度也不会把目标裁掉相对安全。如果要加亮度扰动幅度也要小因为红外图像的灰度本身就携带了目标的强度信息扰动太大会破坏这个特征。import random def augment(img, mask): if random.random() 0.5: img np.flip(img, axis2).copy() mask np.flip(mask, axis2).copy() if random.random() 0.5: img np.flip(img, axis1).copy() mask np.flip(mask, axis1).copy() k random.choice([0, 1, 2, 3]) if k 0: img np.rot90(img, k, axes(1, 2)).copy() mask np.rot90(mask, k, axes(1, 2)).copy() return img, mask注意flip和rot90之后要.copy()否则numpy返回的是视图后续操作可能出问题。这个坑我踩过一次表现为训练时loss莫名其妙不下降排查半天才发现是数据增强返回了共享内存的视图。3.4 训练集和测试集的划分数据量小的时候划分要讲究。我按8:2划分训练和测试并且保证两个集合里都包含各种场景天空、地面、海面。如果随机划分导致测试集全是简单场景评估结果会虚高没有参考价值。划分完之后建议把训练集里目标像素占比统计一下。如果某些图里目标像素占比极低比如低于0.01%训练时这些图的loss贡献几乎为零可以考虑适当过采样或者用带权重的损失函数补偿。4. DNANet网络结构的逐模块实现4.1 整体架构的拆解思路DNANet整体是编码器-解码器结构编码器逐层下采样提取特征解码器逐层上采样恢复分辨率中间用密集嵌套连接把各层特征反复融合。和U-Net的区别在于U-Net的跳跃连接是一对一的而DNANet是多对多的每一层都接收前面所有层的特征。实现的时候我建议先把整体骨架搭出来用占位模块跑通前向传播确认张量维度都对再逐个替换成真正的模块。这样调试成本最低不会一上来就被一堆维度错误淹没。4.2 基础卷积块的设计基础卷积块是网络的积木设计上要兼顾表达能力和计算量。我用的是卷积BN激活的组合激活函数选LeakyReLU而不是ReLU原因是红外小目标的响应本来就弱ReLU在负半轴直接置零容易把弱信号彻底掐掉LeakyReLU保留一点负半轴斜率对小目标更友好。import torch.nn as nn class ConvBlock(nn.Module): def __init__(self, in_ch, out_ch, kernel_size3, stride1): super().__init__() padding kernel_size // 2 self.conv nn.Conv2d(in_ch, out_ch, kernel_size, stride, padding, biasFalse) self.bn nn.BatchNorm2d(out_ch) self.act nn.LeakyReLU(0.1, inplaceTrue) def forward(self, x): return self.act(self.bn(self.conv(x)))bias设为False是因为后面接了BNBN本身有偏置再叠加卷积的bias是冗余的还多占显存。这个细节很多人不注意但在小显存卡上积少成多。4.3 密集嵌套融合模块的实现这是DNANet的核心。密集嵌套的意思是解码器第i层的输入来自编码器对应层和解码器前面所有层输出的拼接。实现上用一个列表缓存各层输出每次融合时把需要的层拼起来。class DenseFusion(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.conv ConvBlock(in_ch, out_ch) def forward(self, feats): x torch.cat(feats, dim1) return self.conv(x)调用的时候把当前层需要的所有特征放进feats列表cat之后过一层卷积降维。这里要注意通道数的计算拼接后的通道数是所有输入通道之和ConvBlock的in_ch必须和这个和一致否则直接报维度错误。我一般会在config里把每层的通道数写清楚避免手算出错。密集嵌套带来的一个副作用是显存占用上升因为要缓存多层特征。如果显存吃紧可以适当减少融合的层数或者用1x1卷积先降维再拼接。4.4 注意力模块CSAM与DNIM的配合注意力部分是DNANet区别于普通密集网络的关键。这里涉及两个模块一个是通道-空间注意力CSAM一个是密集嵌套交互模块DNIM两者配合完成特征的重新加权。CSAM的思路是先做通道注意力再做空间注意力串行叠加。通道注意力回答哪些通道更重要空间注意力回答图上哪些位置更重要。实现上通道注意力用全局平均池化加两层全连接空间注意力用通道维度的最大池化和平均池化拼接后过一层卷积。class CSAM(nn.Module): def __init__(self, channels, reduction16): super().__init__() self.avg_pool nn.AdaptiveAvgPool2d(1) self.fc nn.Sequential( nn.Linear(channels, channels // reduction, biasFalse), nn.ReLU(inplaceTrue), nn.Linear(channels // reduction, channels, biasFalse), nn.Sigmoid() ) self.spatial nn.Conv2d(2, 1, kernel_size7, padding3, biasFalse) self.sigmoid nn.Sigmoid() def forward(self, x): b, c, _, _ x.size() y self.avg_pool(x).view(b, c) y self.fc(y).view(b, c, 1, 1) x x * y avg_out torch.mean(x, dim1, keepdimTrue) max_out, _ torch.max(x, dim1, keepdimTrue) s torch.cat([avg_out, max_out], dim1) s self.sigmoid(self.spatial(s)) return x * sreduction取16是常规做法通道数少的时候可以取8避免降维后信息损失太多。空间注意力的卷积核用7x7感受野够大能覆盖小目标周围的上下文。DNIM则负责在密集嵌套的路径上做交互把不同层级的特征通过注意力加权后再融合。它的实现比CSAM复杂一些核心是把相邻层的特征做交叉注意力让深层语义指导浅层特征的筛选同时浅层细节补充深层的空间信息。注意注意力模块不要堆太多。我一开始在每一层融合后都加了CSAM结果显存爆了不说训练还变得很不稳定。后来改成只在关键的三层融合后加效果反而更好。注意力是双刃剑加多了会过拟合尤其是在小数据集上。4.5 上采样与输出层解码器的上采样用双线性插值而不是转置卷积。转置卷积容易产生棋盘格伪影对小目标检测是致命的因为伪影本身就像小亮点会干扰判断。双线性插值平滑没有这个问题。class UpBlock(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.up nn.Upsample(scale_factor2, modebilinear, align_cornersTrue) self.conv ConvBlock(in_ch, out_ch) def forward(self, x): return self.conv(self.up(x))输出层用1x1卷积把通道数降到1然后过Sigmoid得到0-1的概率图。训练时和掩码算二值交叉熵推理时按阈值二值化。5. 训练流程与损失函数的选择5.1 为什么不能只用BCE红外小目标检测有个天然的类别不平衡问题目标像素可能只占全图的千分之一甚至更少背景像素占绝大多数。如果只用二值交叉熵BCE网络会发现全部预测为背景就能拿到很低的loss于是学不到任何东西这就是所谓的背景主导。解决办法是组合损失。我用的方案是BCE加Dice损失BCE负责逐像素的稳定梯度Dice负责优化整体的重叠度两者互补。Dice对小目标尤其友好因为它关注的是预测和真值的重叠比例而不是绝对像素数。class BCEDiceLoss(nn.Module): def __init__(self, bce_weight0.5): super().__init__() self.bce_weight bce_weight self.bce nn.BCELoss() def forward(self, pred, target): bce_loss self.bce(pred, target) smooth 1e-6 pred_flat pred.view(-1) target_flat target.view(-1) intersection (pred_flat * target_flat).sum() dice_loss 1 - (2 * intersection smooth) / \ (pred_flat.sum() target_flat.sum() smooth) return self.bce_weight * bce_loss \ (1 - self.bce_weight) * dice_lossbce_weight取0.5是个稳妥的起点如果发现召回率偏低漏检多可以调低bce权重让Dice占主导如果误检多就调高bce权重。5.2 优化器和学习率的设置优化器用Adam初始学习率1e-3。Adam对学习率不那么敏感适合快速起步。但训练到后期固定学习率会导致loss在最优值附近震荡所以要用学习率调度。我用的是余弦退火加warmup。warmup在前几个epoch把学习率从很小的值线性升到初始值避免一开始梯度太大把网络带偏余弦退火让学习率平滑下降到接近零帮助收敛到更优的点。from torch.optim.lr_scheduler import CosineAnnealingLR, LambdaLR def get_scheduler(optimizer, warmup_epochs, total_epochs): def lr_lambda(epoch): if epoch warmup_epochs: return (epoch 1) / warmup_epochs progress (epoch - warmup_epochs) / \ (total_epochs - warmup_epochs) return 0.5 * (1 np.cos(np.pi * progress)) return LambdaLR(optimizer, lr_lambda)warmup一般设3-5个epoch总epoch数根据数据量定小数据集100-200个epoch比较合适。5.3 训练循环的关键细节训练循环本身不复杂但有几个细节决定成败。第一每个epoch开始前要调用model.train()结束评估时调用model.eval()并且用torch.no_grad()包住评估否则BN和Dropout的行为不对显存也会白白占用。第二梯度要清零。PyTorch默认累积梯度忘记zero_grad()会导致梯度越积越大训练直接崩。第三如果显存不够用梯度累积模拟大batch。比如想要batch size 16但显存只够4就累积4次梯度再更新一次。for epoch in range(total_epochs): model.train() for img, mask in train_loader: img, mask img.to(device), mask.to(device) pred model(img) loss criterion(pred, mask) optimizer.zero_grad() loss.backward() optimizer.step() scheduler.step() # 评估 model.eval() with torch.no_grad(): for img, mask in test_loader: img, mask img.to(device), mask.to(device) pred model(img) # 计算指标5.4 训练过程中的监控指标光看loss不够loss下降不代表检测效果好。要同时监控几个指标IoU交并比、Pd检测率即召回率、Fa虚警率。IoU衡量整体重叠度Pd衡量漏检情况Fa衡量误检情况。小目标检测里Pd和Fa往往是一对矛盾调参就是在两者之间找平衡。我习惯每个epoch把这三个指标打印出来画成曲线。如果Pd上不去说明网络没学到目标如果Fa太高说明背景杂波没压住可能要加注意力或者调损失权重。6. 推理与后处理让结果真正可用6.1 推理流程和阈值选择推理时把测试图送进网络得到0-1的概率图然后按阈值二值化。阈值的选择很关键太高会漏检太低会误检。默认0.5不一定最优我一般会在验证集上扫一遍阈值从0.1到0.9看哪个阈值下F1最高。def inference(model, img, threshold0.5): model.eval() with torch.no_grad(): img img.to(device) pred model(img) pred (pred threshold).float() return pred如果测试集和训练集分布差异大阈值可能要重新调。这个没有一劳永逸的办法只能根据实际数据来。6.2 连通域后处理去噪网络输出的二值图里往往有一些孤立的噪点这些噪点面积很小明显不是目标。可以用连通域分析过滤掉面积过小的区域。import cv2 import numpy as np def remove_small_regions(mask, min_area3): mask mask.astype(np.uint8) num, labels, stats, _ cv2.connectedComponentsWithStats(mask, 8) out np.zeros_like(mask) for i in range(1, num): if stats[i, cv2.CC_STAT_AREA] min_area: out[labels i] 1 return outmin_area要根据目标尺度定如果目标最小只有几个像素min_area设太大反而会把真目标滤掉。我一般设3-5具体看数据集。6.3 可视化验证推理完一定要可视化把原图、真值掩码、预测掩码并排画出来。光看指标数字容易自欺欺人可视化能一眼看出问题是漏检了、误检了还是边界对不齐。import matplotlib.pyplot as plt def visualize(img, gt, pred): fig, axes plt.subplots(1, 3, figsize(12, 4)) axes[0].imshow(img[0], cmapgray) axes[0].set_title(Image) axes[1].imshow(gt[0], cmapgray) axes[1].set_title(Ground Truth) axes[2].imshow(pred[0], cmapgray) axes[2].set_title(Prediction) for ax in axes: ax.axis(off) plt.show()我习惯每次评估都随机抽几张图看尤其是那些指标差的图往往能发现系统性问题。7. 复现过程中踩过的坑与调参经验7.1 小目标在深层特征图消失这是复现初期最典型的问题。训练了几十个epochloss降下去了但预测图里小目标几乎全丢。排查下来发现是下采样太激进目标在深层特征图上已经不足一个像素。解决办法有两个一是减少下采样次数把编码器从5层减到4层二是在密集嵌套里强化浅层特征的权重让高分辨率特征更多地参与深层计算。我两个都用了效果明显改善。7.2 损失不下降的几种原因loss不下降的原因很多我按排查顺序列一下先看数据加载对不对把batch里的图和掩码可视化出来确认配对没错再看学习率是不是太大太大导致震荡太小导致不动然后看损失函数实现有没有bug尤其是Dice的分母有没有加smooth防止除零最后看网络输出有没有经过Sigmoid如果输出是logits却直接喂给BCELossloss会异常。7.3 过拟合的识别与缓解小数据集上过拟合几乎是必然的。识别方法很简单训练loss持续下降但验证指标到某个点后开始变差。缓解手段按优先级排先加数据增强再加Dropout或权重衰减最后考虑减小模型容量。我在全连接层前加了Dropout(0.3)并在优化器里设了weight_decay1e-4过拟合明显缓解。但要注意Dropout不要加在卷积层后面会破坏空间信息。7.4 显存不足的应对策略显存不足时按这个顺序优化先减小batch size这是最直接的再检查有没有不必要的中间变量占用显存比如在训练循环里保留了计算图然后考虑混合精度训练用torch.cuda.amp能省将近一半显存最后才是减小模型。混合精度训练的实现from torch.cuda.amp import autocast, GradScaler scaler GradScaler() for img, mask in train_loader: optimizer.zero_grad() with autocast(): pred model(img) loss criterion(pred, mask) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()混合精度要注意某些操作在fp16下会溢出如果发现loss变成nan先关掉amp排查。7.5 指标波动的正常范围训练过程中指标有波动是正常的尤其是小数据集。如果Pd在几个百分点内波动不用太在意如果波动超过10个百分点说明训练不稳定要检查学习率、batch size或者数据增强的随机性。我一般会跑多次实验取平均单次结果参考价值有限。如果时间紧至少固定随机种子跑一次保证结果可复现。8. 从复现到改进几个可以尝试的方向8.1 骨干网络的替换DNANet原版用的骨干比较轻量如果想提升性能可以换成更强的骨干比如ResNet系列。但要注意骨干越深小目标消失的风险越大所以换骨干的同时要强化密集嵌套的融合或者引入多尺度特征。8.2 损失函数的进一步优化BCE加Dice是基础组合还可以尝试Focal Loss或者Tversky Loss。Focal Loss通过调制因子让网络更关注难样本对小目标有帮助Tversky Loss可以调节假阳和假阴的权重适合对漏检或误检有偏好的场景。8.3 多尺度训练与测试小目标的尺度变化是个难题。多尺度训练指在训练时随机缩放输入让网络见到不同尺度的目标多尺度测试指在推理时对同一张图用多个尺度预测再融合。这两个技巧都能提升尺度鲁棒性代价是训练和推理时间增加。8.4 部署时的优化考虑如果最终要部署推理速度是个硬指标。可以做的优化包括把BN层折叠进卷积、用TensorRT或ONNX Runtime加速、量化到INT8。这些优化对精度有轻微影响需要重新评估指标。我在实际部署时发现BN折叠几乎无损量化到INT8会掉一两个点的IoU如果对精度要求高建议只做BN折叠和算子融合。整套流程走下来DNANet的复现难点其实不在网络结构本身而在数据质量、损失设计和训练细节的把控上。网络结构照着论文搭维度对齐了就能跑但要让小目标真正不丢得在数据增强、损失权重、注意力模块的取舍上反复试。我个人的体会是先把baseline跑通指标能到论文的八九成再逐步加改进比一上来就堆模块要稳得多。另外红外小目标这个方向对数据的依赖极强如果手头数据场景单一建议先扩充数据或者做场景迁移否则调参调到天荒地老也难有突破。