：残差+注意力+软阈值化的抗噪分类实战)
深度残差收缩网络Deep Residual Shrinkage NetworkDRSN这名字听着唬人实际拆开就是三件事残差结构、注意力机制和软阈值化。最近我把它的完整 Python 程序从头到尾复现了一遍专门测试了它在带噪声图像分类上的表现。这篇文章会把 DRSN 的来龙去脉、核心原理、完整训练代码和踩坑记录都摊开来讲适合已经会 PyTorch 基础操作、想搞懂残差网络变体或者想拿深度学习方法做信号/图像降噪的读者。如果你只需要跑通一个程序可以直接跳到第 2 节的代码如果你想知道为什么这个网络有效建议顺着读。我用的环境是 Python 3.10 PyTorch 2.1没有特别花哨的依赖CPU 也能完成小规模训练。1. 深度残差收缩网络到底在解决什么问题1.1 普通残差网络的局限噪声会进入特征ResNet 通过恒等映射解决了深层网络退化问题让梯度可以顺畅回传这是它成功的关键。但有个点很容易被忽略ResNet 对特征图的每个通道是同等对待的卷积核会把输入里的所有响应都编码进特征包括噪声。当数据来自传感器、工业现场或低光照环境时原始信号里会混入大量干扰模型训练时如果把这些噪声也当成有效模式测试时大概率先下降。传统做法是先滤波再分类。比如振动信号先用小波去噪图像先用双边滤波但这类预处理会损失边缘和细节而且滤波参数需要人工调换一个工况又要重调。DRSN 的思路完全不同它把降噪过程直接嵌入网络让网络根据当前样本自己决定哪些特征需要保留、哪些特征该被压缩。这样模型既保留残差结构带来的训练稳定性又能对噪声产生天然的抑制能力。1.2 软阈值化怎么做到“部分去噪”软阈值化的公式是 y sign(x) · max(|x| − τ, 0)。直观解释就是在数轴上把绝对值小于 τ 的数直接变成 0绝对值大于 τ 的数向 0 方向收缩 τ。这比硬阈值化绝对值小于 τ 置 0大于 τ 不变更温和避免了输出在阈值处突然跳变梯度也更稳定。我用一个生活化类比你在收音机里听广播背景噪声小的时候不需要处理噪声大的时候把音量整体压低一些这样人声虽然也变轻但不会淹没在杂音里。软阈值化就是这种“压低”操作τ 是压低幅度。常规降噪算法需要预先知道噪声范围而 DRSN 的 τ 不是人工定的是由网络自己学出来的。1.3 注意力机制决定阈值大小如果每个通道用同一个固定阈值显然不够聪明。不同通道对最终分类的贡献不一样包含噪声的比例也不一样。DRSN 借鉴通道注意力机制用一个极小的子网络生成阈值先对特征图做全局平均池化得到每个通道的整体响应再经过两个全连接层中间用 ReLU 激活最后一层用 Sigmoid 把输出压缩到 0 到 1 之间最后用这个缩放系数乘以特征图绝对值的均值得到真正的阈值。为什么要用 Sigmoid 乘以均值因为 Sigmoid 输出永远不会超过 1所以阈值一定小于当前特征图的平均绝对值。如果阈值比所有特征值都大软阈值化会把整张特征图清零信息全部丢失如果太小又起不到去噪效果。乘均值这个操作保证了阈值和输入尺度匹配同时保留了一定的学习空间。这个设计是 DRSN 能稳定训练的核心原因。2. 环境准备与模型结构拆解2.1 依赖安装与工程目录复现 DRSN 不需要太多第三方库。我的环境只需要 PyTorch 和 torchvision画曲线再加一个 matplotlib。安装命令如下pip install torch torchvision matplotlib为了让你快速跑起来我把所有代码放在一个drsn_train.py文件里不拆模块。单文件的优点是直接复制就能运行不用建一堆目录缺点是代码稍长但胜在清晰。文件内部按这个顺序组织导入库、定义阈值生成器、定义残差收缩块、定义 DRSN 主体、加载 CIFAR-10、训练循环、测试循环。2.2 阈值生成器DRSN 的核心小网络ThresholdGenerator 是整个模型的灵魂。它的输入是某一个残差块的输出特征图输出是每个通道对应的阈值。代码如下import torch import torch.nn as nn import torch.nn.functional as F class ThresholdGenerator(nn.Module): def __init__(self, in_channels, reduction16, modechannel_wise): super().__init__() self.mode mode self.avg_pool nn.AdaptiveAvgPool2d(1) self.fc1 nn.Linear(in_channels, in_channels // reduction) self.fc2 nn.Linear(in_channels // reduction, in_channels if mode channel_wise else 1) self.sigmoid nn.Sigmoid() def forward(self, x): avg_abs torch.mean(torch.abs(x), dim(2, 3), keepdimTrue) gap self.avg_pool(x).view(x.size(0), -1) scale self.fc2(F.relu(self.fc1(gap))) scale self.sigmoid(scale).view(-1, scale.size(1), 1, 1) threshold avg_abs * scale return threshold逐行解释一下。avg_abs是在每个通道的空间维度上求绝对值均值形状是[N, C, 1, 1]。gap是常规的全局平均池化先通过AdaptiveAvgPool2d(1)把特征图压成[N, C, 1, 1]再view成[N, C]。经过两层全连接后scale的形状是[N, C]channel_wise 模式或[N, 1]channel_shared 模式。Sigmoid 把它压到 0~1再还原成 1×1 的卷积形状方便和avg_abs做逐通道乘法。这里的reduction16是通道压缩比率可以理解为两个全连接之间的瓶颈宽度。缩小reduction会增加参数量但也可能提升拟合能力我试过reduction8和reduction16在 CIFAR-10 上差距不大16 更省参数。2.3 残差收缩块把软阈值化放进残差路径有了阈值生成器接下来把它放进残差块里。DRSN 的残差块和普通 ResNet BasicBlock 的区别在于普通残差块是“卷积-批归一化-激活-卷积-批归一化-恒等映射相加-激活”DRSN 在最后一个批归一化之后、恒等映射相加之前插入了软阈值化操作。代码实现如下class BasicShrinkageBlock(nn.Module): expansion 1 def __init__(self, in_channels, out_channels, stride1, downsampleNone, reduction16, modechannel_wise): super().__init__() self.conv1 nn.Conv2d(in_channels, out_channels, kernel_size3, stridestride, padding1, biasFalse) self.bn1 nn.BatchNorm2d(out_channels) self.relu nn.ReLU(inplaceTrue) self.conv2 nn.Conv2d(out_channels, out_channels, kernel_size3, stride1, padding1, biasFalse) self.bn2 nn.BatchNorm2d(out_channels) self.threshold_generator ThresholdGenerator(out_channels, reduction, mode) self.downsample downsample def _soft_threshold(self, x, threshold): return torch.sign(x) * torch.clamp(torch.abs(x) - threshold, min0) def forward(self, x): identity x out self.conv1(x) out self.bn1(out) out self.relu(out) out self.conv2(out) out self.bn2(out) threshold self.threshold_generator(out) out self._soft_threshold(out, threshold) if self.downsample is not None: identity self.downsample(x) out identity out self.relu(out) return out可以看到软阈值化只作用在残差分支上不直接作用在恒等映射上。这样做的好处是如果阈值把残留分支的大部分信号压掉恒等映射仍然保留原始输入梯度依然能顺畅回传。downsample用于改变通道数或降低分辨率通常是一个 1×1 卷积加批归一化。2.4 组装完整 DRSN 网络把上述残差块堆叠起来就得到完整的 DRSN。我的配置针对 32×32 小图像做了调整没有照搬 ImageNet 版本class DRSNet(nn.Module): def __init__(self, blockBasicShrinkageBlock, layers(2, 2, 2), channels(16, 32, 64), num_classes10, modechannel_wise): super().__init__() self.in_channels 16 self.conv1 nn.Conv2d(3, 16, kernel_size3, stride1, padding1, biasFalse) self.bn1 nn.BatchNorm2d(16) self.relu nn.ReLU(inplaceTrue) self.layer1 self._make_layer(block, channels[0], layers[0], stride1, modemode) self.layer2 self._make_layer(block, channels[1], layers[1], stride2, modemode) self.layer3 self._make_layer(block, channels[2], layers[2], stride2, modemode) self.avgpool nn.AdaptiveAvgPool2d((1, 1)) self.fc nn.Linear(channels[2] * block.expansion, num_classes) def _make_layer(self, block, out_channels, blocks, stride1, modechannel_wise): downsample None if stride ! 1 or self.in_channels ! out_channels: downsample nn.Sequential( nn.Conv2d(self.in_channels, out_channels, kernel_size1, stridestride, biasFalse), nn.BatchNorm2d(out_channels), ) layers [] layers.append(block(self.in_channels, out_channels, stridestride, downsampledownsample, modemode)) self.in_channels out_channels for _ in range(1, blocks): layers.append(block(self.in_channels, out_channels, modemode)) return nn.Sequential(*layers) def forward(self, x): x self.conv1(x) x self.bn1(x) x self.relu(x) x self.layer1(x) x self.layer2(x) x self.layer3(x) x self.avgpool(x) x torch.flatten(x, 1) x self.fc(x) return x这个配置下输入 32×32 图像经过第一个卷积后仍为 32×32layer1 输出 32×32layer2 输出 16×16layer3 输出 8×8最后平均池化得到 1×1 特征。如果想把 DRSN 用到更大的输入只需要把layers和channels加长。3. 完整 Python 训练程序与实操3.1 数据准备手动给 CIFAR-10 加噪声我用 CIFAR-10 做演示因为它是公开数据集加载方便而且标签是 10 分类训练速度快。为了突出 DRSN 的抗噪能力我在训练和测试时都往图像上叠加高斯噪声。数据增强部分用常规的随机裁剪和水平翻转import torchvision import torchvision.transforms as transforms transform_train transforms.Compose([ transforms.RandomCrop(32, padding4), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), ]) transform_test transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), ]) train_set torchvision.datasets.CIFAR10(root./data, trainTrue, downloadTrue, transformtransform_train) test_set torchvision.datasets.CIFAR10(root./data, trainFalse, downloadTrue, transformtransform_test)加噪声函数写成一个普通 Python 函数在训练循环里每个 batch 调用。噪声标准差我用 0.1这个值对归一化后的图像来说已经不小。如果设为 0.3普通 ResNet 基本会崩DRSN 也会明显下降如果设为 0.01又体现不出差距。def add_noise(images, noise_std0.1): noise torch.randn_like(images) * noise_std return torch.clamp(images noise, -1, 1)注意Normalize((0.5,0.5,0.5),(0.5,0.5,0.5))后图像的数值范围是 [-1,1]噪声也要在这个量级加最后用clamp防止越界。3.2 训练循环超参数与学习率调整训练部分的核心代码示例如下。这里我用 Adam学习率 0.001训练 20 个 epoch每 10 个 epoch 把学习率降到原来的十分之一。为了复现稳定先固定随机种子import torch.optim as optim import numpy as np import random def set_seed(seed42): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) set_seed(42) device torch.device(cuda if torch.cuda.is_available() else cpu) model DRSNet().to(device) criterion nn.CrossEntropyLoss() optimizer optim.Adam(model.parameters(), lr1e-3, weight_decay5e-4) scheduler optim.lr_scheduler.MultiStepLR(optimizer, milestones[10, 15], gamma0.1) train_loader torch.utils.data.DataLoader(train_set, batch_size128, shuffleTrue, num_workers2) test_loader torch.utils.data.DataLoader(test_set, batch_size128, shuffleFalse, num_workers2)训练循环里最值得注意的就是加噪声的时机。我是在数据加载完成之后、送入网络之前加噪声这样每个 epoch 看到的噪声都不同相当于一种数据增强loss_history [] for epoch in range(20): model.train() running_loss 0.0 for images, labels in train_loader: images add_noise(images.to(device), noise_std0.1).to(device) labels labels.to(device) optimizer.zero_grad() outputs model(images) loss criterion(outputs, labels) loss.backward() optimizer.step() running_loss loss.item() * images.size(0) epoch_loss running_loss / len(train_set) loss_history.append(epoch_loss) print(fEpoch {epoch1:02d} | Loss: {epoch_loss:.4f}) scheduler.step()为什么要在 batch 上加噪声而不是在数据集里加如果事先把整个训练集加一遍噪声模型只会学习固定的噪声模式过几个 epoch 就记住了。动态加噪可以让模型见过更多噪声变化泛化能力更强。3.3 评估与结果可视化测试阶段同样加噪声但噪声强度和训练时保持一致。为了公平对比我每次评估用一个固定的随机种子生成噪声这样多次跑结果不会飘得太离谱。def evaluate(loader, model, noise_std0.1): model.eval() correct 0 total 0 with torch.no_grad(): for images, labels in loader: images add_noise(images.to(device), noise_stdnoise_std).to(device) labels labels.to(device) outputs model(images) preds outputs.argmax(dim1) correct (preds labels).sum().item() total labels.size(0) return correct / total test_acc evaluate(test_loader, model, noise_std0.1) print(fTest accuracy with noise: {test_acc:.4f})如果想把训练损失画出来用 matplotlib 画个折线图就行。我实际复现时发现DRSN 在加载噪声后的前几个 epoch loss 下降很慢但随着阈值生成器学会控制阈值后面会越来越快。4. 复现中的几个大坑建议收藏4.1 阈值输出过大导致特征全部缩死我第一次实现时让阈值生成器的最后一层直接输出阈值结果训练不到几个 batch loss 就变成 NaN。原因很简单全连接层输出的数值范围不受控制可能产生非常大的阈值软阈值化把所有特征都压成 0。一旦特征全为 0残差分支没有有效信号模型只剩下恒等映射和最后的分类层梯度信号不足训练直接崩溃。解决办法就是前面代码里的方案最后一层用 Sigmoid 输出 0~1 的缩放系数再乘以特征图绝对值均值。这样阈值永远落在 0 到当前通道平均绝对值之间不会出现极端值。这个改动让训练从“随缘”变成了“稳”。4.2 软阈值化求导要小心PyTorch 的torch.sign在 0 点不可导通常会被认为梯度为 0。但在torch.sign(x) * torch.clamp(torch.abs(x) - threshold, min0)这个表达式里sign是作为常数系数参与乘法的。梯度的主要来源是abs(x)和clamp部分所以模型依然能正常回传。如果你还是不放心可以换一个等价写法def soft_threshold(x, threshold): return torch.sign(x) * torch.relu(torch.abs(x) - threshold)我更推荐把这个操作封装成nn.Module方便复用。4.3 小图像上别照搬大网络结构一些网上公开的实现直接使用 ImageNet 版本的 ResNet 结构第一层就是 7×7 卷积加 stride2输入 32×32 时很快变成 16×16信息损失非常严重。我改成 3×3、stride1 的起始卷积并让 layer1 保持分辨率不变。如果你要处理的是 64×64 或更大图像可以适当增加 stage 数量。4.4 参数量与训练速度的权衡DRSN 比普通 ResNet 多的部分主要是阈值生成器的两个全连接层参数量增加很少。我按现在的配置统计过普通 ResNet-10 大概 1.2M 参数DRSN-CW-10 大概 1.3M多出的 0.1M 基本就是每个残差块里那两层全连接。换来的是带噪场景下测试准确率的明显提升。如果对推理速度敏感可以把mode改成channel_shared这样每次只生成一个阈值计算量更小但精度会略低于channel_wise。5. 最后再分享一点我的实测体会如果你在自己的数据集上测试我建议先用小噪声强度跑通再用大噪声强度看 DRSN 的收益。我跑下来的现象是干净数据上 DRSN 并不比普通 ResNet 强甚至偶尔低零点几个百分点但只要测试时加了噪声两者的差距立刻拉开。这基本符合 DRSN 的设计目标——它本来就不是为了提升干净数据的精度而是为了让模型在噪声干扰下不崩掉。还有个细节训练时加噪声强度不要一直固定可以先小后大比如前 10 个 epoch 用 0.05后 10 个 epoch 用 0.1收敛速度会更快。希望这份复现记录能帮你少踩几个坑。
深度残差收缩网络Deep Residual Shrinkage NetworkDRSN这名字听着唬人实际拆开就是三件事残差结构、注意力机制和软阈值化。最近我把它的完整 Python 程序从头到尾复现了一遍专门测试了它在带噪声图像分类上的表现。这篇文章会把 DRSN 的来龙去脉、核心原理、完整训练代码和踩坑记录都摊开来讲适合已经会 PyTorch 基础操作、想搞懂残差网络变体或者想拿深度学习方法做信号/图像降噪的读者。如果你只需要跑通一个程序可以直接跳到第 2 节的代码如果你想知道为什么这个网络有效建议顺着读。我用的环境是 Python 3.10 PyTorch 2.1没有特别花哨的依赖CPU 也能完成小规模训练。1. 深度残差收缩网络到底在解决什么问题1.1 普通残差网络的局限噪声会进入特征ResNet 通过恒等映射解决了深层网络退化问题让梯度可以顺畅回传这是它成功的关键。但有个点很容易被忽略ResNet 对特征图的每个通道是同等对待的卷积核会把输入里的所有响应都编码进特征包括噪声。当数据来自传感器、工业现场或低光照环境时原始信号里会混入大量干扰模型训练时如果把这些噪声也当成有效模式测试时大概率先下降。传统做法是先滤波再分类。比如振动信号先用小波去噪图像先用双边滤波但这类预处理会损失边缘和细节而且滤波参数需要人工调换一个工况又要重调。DRSN 的思路完全不同它把降噪过程直接嵌入网络让网络根据当前样本自己决定哪些特征需要保留、哪些特征该被压缩。这样模型既保留残差结构带来的训练稳定性又能对噪声产生天然的抑制能力。1.2 软阈值化怎么做到“部分去噪”软阈值化的公式是 y sign(x) · max(|x| − τ, 0)。直观解释就是在数轴上把绝对值小于 τ 的数直接变成 0绝对值大于 τ 的数向 0 方向收缩 τ。这比硬阈值化绝对值小于 τ 置 0大于 τ 不变更温和避免了输出在阈值处突然跳变梯度也更稳定。我用一个生活化类比你在收音机里听广播背景噪声小的时候不需要处理噪声大的时候把音量整体压低一些这样人声虽然也变轻但不会淹没在杂音里。软阈值化就是这种“压低”操作τ 是压低幅度。常规降噪算法需要预先知道噪声范围而 DRSN 的 τ 不是人工定的是由网络自己学出来的。1.3 注意力机制决定阈值大小如果每个通道用同一个固定阈值显然不够聪明。不同通道对最终分类的贡献不一样包含噪声的比例也不一样。DRSN 借鉴通道注意力机制用一个极小的子网络生成阈值先对特征图做全局平均池化得到每个通道的整体响应再经过两个全连接层中间用 ReLU 激活最后一层用 Sigmoid 把输出压缩到 0 到 1 之间最后用这个缩放系数乘以特征图绝对值的均值得到真正的阈值。为什么要用 Sigmoid 乘以均值因为 Sigmoid 输出永远不会超过 1所以阈值一定小于当前特征图的平均绝对值。如果阈值比所有特征值都大软阈值化会把整张特征图清零信息全部丢失如果太小又起不到去噪效果。乘均值这个操作保证了阈值和输入尺度匹配同时保留了一定的学习空间。这个设计是 DRSN 能稳定训练的核心原因。2. 环境准备与模型结构拆解2.1 依赖安装与工程目录复现 DRSN 不需要太多第三方库。我的环境只需要 PyTorch 和 torchvision画曲线再加一个 matplotlib。安装命令如下pip install torch torchvision matplotlib为了让你快速跑起来我把所有代码放在一个drsn_train.py文件里不拆模块。单文件的优点是直接复制就能运行不用建一堆目录缺点是代码稍长但胜在清晰。文件内部按这个顺序组织导入库、定义阈值生成器、定义残差收缩块、定义 DRSN 主体、加载 CIFAR-10、训练循环、测试循环。2.2 阈值生成器DRSN 的核心小网络ThresholdGenerator 是整个模型的灵魂。它的输入是某一个残差块的输出特征图输出是每个通道对应的阈值。代码如下import torch import torch.nn as nn import torch.nn.functional as F class ThresholdGenerator(nn.Module): def __init__(self, in_channels, reduction16, modechannel_wise): super().__init__() self.mode mode self.avg_pool nn.AdaptiveAvgPool2d(1) self.fc1 nn.Linear(in_channels, in_channels // reduction) self.fc2 nn.Linear(in_channels // reduction, in_channels if mode channel_wise else 1) self.sigmoid nn.Sigmoid() def forward(self, x): avg_abs torch.mean(torch.abs(x), dim(2, 3), keepdimTrue) gap self.avg_pool(x).view(x.size(0), -1) scale self.fc2(F.relu(self.fc1(gap))) scale self.sigmoid(scale).view(-1, scale.size(1), 1, 1) threshold avg_abs * scale return threshold逐行解释一下。avg_abs是在每个通道的空间维度上求绝对值均值形状是[N, C, 1, 1]。gap是常规的全局平均池化先通过AdaptiveAvgPool2d(1)把特征图压成[N, C, 1, 1]再view成[N, C]。经过两层全连接后scale的形状是[N, C]channel_wise 模式或[N, 1]channel_shared 模式。Sigmoid 把它压到 0~1再还原成 1×1 的卷积形状方便和avg_abs做逐通道乘法。这里的reduction16是通道压缩比率可以理解为两个全连接之间的瓶颈宽度。缩小reduction会增加参数量但也可能提升拟合能力我试过reduction8和reduction16在 CIFAR-10 上差距不大16 更省参数。2.3 残差收缩块把软阈值化放进残差路径有了阈值生成器接下来把它放进残差块里。DRSN 的残差块和普通 ResNet BasicBlock 的区别在于普通残差块是“卷积-批归一化-激活-卷积-批归一化-恒等映射相加-激活”DRSN 在最后一个批归一化之后、恒等映射相加之前插入了软阈值化操作。代码实现如下class BasicShrinkageBlock(nn.Module): expansion 1 def __init__(self, in_channels, out_channels, stride1, downsampleNone, reduction16, modechannel_wise): super().__init__() self.conv1 nn.Conv2d(in_channels, out_channels, kernel_size3, stridestride, padding1, biasFalse) self.bn1 nn.BatchNorm2d(out_channels) self.relu nn.ReLU(inplaceTrue) self.conv2 nn.Conv2d(out_channels, out_channels, kernel_size3, stride1, padding1, biasFalse) self.bn2 nn.BatchNorm2d(out_channels) self.threshold_generator ThresholdGenerator(out_channels, reduction, mode) self.downsample downsample def _soft_threshold(self, x, threshold): return torch.sign(x) * torch.clamp(torch.abs(x) - threshold, min0) def forward(self, x): identity x out self.conv1(x) out self.bn1(out) out self.relu(out) out self.conv2(out) out self.bn2(out) threshold self.threshold_generator(out) out self._soft_threshold(out, threshold) if self.downsample is not None: identity self.downsample(x) out identity out self.relu(out) return out可以看到软阈值化只作用在残差分支上不直接作用在恒等映射上。这样做的好处是如果阈值把残留分支的大部分信号压掉恒等映射仍然保留原始输入梯度依然能顺畅回传。downsample用于改变通道数或降低分辨率通常是一个 1×1 卷积加批归一化。2.4 组装完整 DRSN 网络把上述残差块堆叠起来就得到完整的 DRSN。我的配置针对 32×32 小图像做了调整没有照搬 ImageNet 版本class DRSNet(nn.Module): def __init__(self, blockBasicShrinkageBlock, layers(2, 2, 2), channels(16, 32, 64), num_classes10, modechannel_wise): super().__init__() self.in_channels 16 self.conv1 nn.Conv2d(3, 16, kernel_size3, stride1, padding1, biasFalse) self.bn1 nn.BatchNorm2d(16) self.relu nn.ReLU(inplaceTrue) self.layer1 self._make_layer(block, channels[0], layers[0], stride1, modemode) self.layer2 self._make_layer(block, channels[1], layers[1], stride2, modemode) self.layer3 self._make_layer(block, channels[2], layers[2], stride2, modemode) self.avgpool nn.AdaptiveAvgPool2d((1, 1)) self.fc nn.Linear(channels[2] * block.expansion, num_classes) def _make_layer(self, block, out_channels, blocks, stride1, modechannel_wise): downsample None if stride ! 1 or self.in_channels ! out_channels: downsample nn.Sequential( nn.Conv2d(self.in_channels, out_channels, kernel_size1, stridestride, biasFalse), nn.BatchNorm2d(out_channels), ) layers [] layers.append(block(self.in_channels, out_channels, stridestride, downsampledownsample, modemode)) self.in_channels out_channels for _ in range(1, blocks): layers.append(block(self.in_channels, out_channels, modemode)) return nn.Sequential(*layers) def forward(self, x): x self.conv1(x) x self.bn1(x) x self.relu(x) x self.layer1(x) x self.layer2(x) x self.layer3(x) x self.avgpool(x) x torch.flatten(x, 1) x self.fc(x) return x这个配置下输入 32×32 图像经过第一个卷积后仍为 32×32layer1 输出 32×32layer2 输出 16×16layer3 输出 8×8最后平均池化得到 1×1 特征。如果想把 DRSN 用到更大的输入只需要把layers和channels加长。3. 完整 Python 训练程序与实操3.1 数据准备手动给 CIFAR-10 加噪声我用 CIFAR-10 做演示因为它是公开数据集加载方便而且标签是 10 分类训练速度快。为了突出 DRSN 的抗噪能力我在训练和测试时都往图像上叠加高斯噪声。数据增强部分用常规的随机裁剪和水平翻转import torchvision import torchvision.transforms as transforms transform_train transforms.Compose([ transforms.RandomCrop(32, padding4), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), ]) transform_test transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), ]) train_set torchvision.datasets.CIFAR10(root./data, trainTrue, downloadTrue, transformtransform_train) test_set torchvision.datasets.CIFAR10(root./data, trainFalse, downloadTrue, transformtransform_test)加噪声函数写成一个普通 Python 函数在训练循环里每个 batch 调用。噪声标准差我用 0.1这个值对归一化后的图像来说已经不小。如果设为 0.3普通 ResNet 基本会崩DRSN 也会明显下降如果设为 0.01又体现不出差距。def add_noise(images, noise_std0.1): noise torch.randn_like(images) * noise_std return torch.clamp(images noise, -1, 1)注意Normalize((0.5,0.5,0.5),(0.5,0.5,0.5))后图像的数值范围是 [-1,1]噪声也要在这个量级加最后用clamp防止越界。3.2 训练循环超参数与学习率调整训练部分的核心代码示例如下。这里我用 Adam学习率 0.001训练 20 个 epoch每 10 个 epoch 把学习率降到原来的十分之一。为了复现稳定先固定随机种子import torch.optim as optim import numpy as np import random def set_seed(seed42): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) set_seed(42) device torch.device(cuda if torch.cuda.is_available() else cpu) model DRSNet().to(device) criterion nn.CrossEntropyLoss() optimizer optim.Adam(model.parameters(), lr1e-3, weight_decay5e-4) scheduler optim.lr_scheduler.MultiStepLR(optimizer, milestones[10, 15], gamma0.1) train_loader torch.utils.data.DataLoader(train_set, batch_size128, shuffleTrue, num_workers2) test_loader torch.utils.data.DataLoader(test_set, batch_size128, shuffleFalse, num_workers2)训练循环里最值得注意的就是加噪声的时机。我是在数据加载完成之后、送入网络之前加噪声这样每个 epoch 看到的噪声都不同相当于一种数据增强loss_history [] for epoch in range(20): model.train() running_loss 0.0 for images, labels in train_loader: images add_noise(images.to(device), noise_std0.1).to(device) labels labels.to(device) optimizer.zero_grad() outputs model(images) loss criterion(outputs, labels) loss.backward() optimizer.step() running_loss loss.item() * images.size(0) epoch_loss running_loss / len(train_set) loss_history.append(epoch_loss) print(fEpoch {epoch1:02d} | Loss: {epoch_loss:.4f}) scheduler.step()为什么要在 batch 上加噪声而不是在数据集里加如果事先把整个训练集加一遍噪声模型只会学习固定的噪声模式过几个 epoch 就记住了。动态加噪可以让模型见过更多噪声变化泛化能力更强。3.3 评估与结果可视化测试阶段同样加噪声但噪声强度和训练时保持一致。为了公平对比我每次评估用一个固定的随机种子生成噪声这样多次跑结果不会飘得太离谱。def evaluate(loader, model, noise_std0.1): model.eval() correct 0 total 0 with torch.no_grad(): for images, labels in loader: images add_noise(images.to(device), noise_stdnoise_std).to(device) labels labels.to(device) outputs model(images) preds outputs.argmax(dim1) correct (preds labels).sum().item() total labels.size(0) return correct / total test_acc evaluate(test_loader, model, noise_std0.1) print(fTest accuracy with noise: {test_acc:.4f})如果想把训练损失画出来用 matplotlib 画个折线图就行。我实际复现时发现DRSN 在加载噪声后的前几个 epoch loss 下降很慢但随着阈值生成器学会控制阈值后面会越来越快。4. 复现中的几个大坑建议收藏4.1 阈值输出过大导致特征全部缩死我第一次实现时让阈值生成器的最后一层直接输出阈值结果训练不到几个 batch loss 就变成 NaN。原因很简单全连接层输出的数值范围不受控制可能产生非常大的阈值软阈值化把所有特征都压成 0。一旦特征全为 0残差分支没有有效信号模型只剩下恒等映射和最后的分类层梯度信号不足训练直接崩溃。解决办法就是前面代码里的方案最后一层用 Sigmoid 输出 0~1 的缩放系数再乘以特征图绝对值均值。这样阈值永远落在 0 到当前通道平均绝对值之间不会出现极端值。这个改动让训练从“随缘”变成了“稳”。4.2 软阈值化求导要小心PyTorch 的torch.sign在 0 点不可导通常会被认为梯度为 0。但在torch.sign(x) * torch.clamp(torch.abs(x) - threshold, min0)这个表达式里sign是作为常数系数参与乘法的。梯度的主要来源是abs(x)和clamp部分所以模型依然能正常回传。如果你还是不放心可以换一个等价写法def soft_threshold(x, threshold): return torch.sign(x) * torch.relu(torch.abs(x) - threshold)我更推荐把这个操作封装成nn.Module方便复用。4.3 小图像上别照搬大网络结构一些网上公开的实现直接使用 ImageNet 版本的 ResNet 结构第一层就是 7×7 卷积加 stride2输入 32×32 时很快变成 16×16信息损失非常严重。我改成 3×3、stride1 的起始卷积并让 layer1 保持分辨率不变。如果你要处理的是 64×64 或更大图像可以适当增加 stage 数量。4.4 参数量与训练速度的权衡DRSN 比普通 ResNet 多的部分主要是阈值生成器的两个全连接层参数量增加很少。我按现在的配置统计过普通 ResNet-10 大概 1.2M 参数DRSN-CW-10 大概 1.3M多出的 0.1M 基本就是每个残差块里那两层全连接。换来的是带噪场景下测试准确率的明显提升。如果对推理速度敏感可以把mode改成channel_shared这样每次只生成一个阈值计算量更小但精度会略低于channel_wise。5. 最后再分享一点我的实测体会如果你在自己的数据集上测试我建议先用小噪声强度跑通再用大噪声强度看 DRSN 的收益。我跑下来的现象是干净数据上 DRSN 并不比普通 ResNet 强甚至偶尔低零点几个百分点但只要测试时加了噪声两者的差距立刻拉开。这基本符合 DRSN 的设计目标——它本来就不是为了提升干净数据的精度而是为了让模型在噪声干扰下不崩掉。还有个细节训练时加噪声强度不要一直固定可以先小后大比如前 10 个 epoch 用 0.05后 10 个 epoch 用 0.1收敛速度会更快。希望这份复现记录能帮你少踩几个坑。
平方复杂度无法被常规优化消除｜TaoToken 统一 Key 通道实测)
