HRNet的PyTorch实现：从高分辨率并行架构到姿态估计实战

新手必看HRNet的PyTorch实现从原理到代码一步步拆解做计算机视觉的人应该对HRNet不陌生这个网络在姿态估计、语义分割、目标检测这些任务上常年霸榜尤其是它那个“高分辨率特征一直保持”的设计思路和大多数先下采样再上采样的Encoder-Decoder结构完全不同。我最早接触HRNet是在做人体关键点检测的时候当时对比过ResNet反卷积、Hourglass这些经典方案最后换成HRNet之后精度提升明显而且训练收敛也更快。这篇文章就把我实际复现HRNet的过程完整记录下来从网络结构拆解到PyTorch代码实现再到训练调参的经验坑点一次性讲清楚。如果你正准备在自己的项目里用HRNet或者想理解为什么它能在高分辨率特征提取上这么能打这篇文章适合你。我会先解释HRNet的核心设计逻辑再给出一个可以跑的PyTorch实现最后附上我在实际训练中踩过的坑和排查思路。不管你是想做姿态估计、分割还是检测的骨干网络替换这篇文章都能给你一个可直接上手的参考。1. HRNet整体设计与思路拆解1.1 为什么HRNet要一直保持高分辨率先看传统做法。像ResNet这类网络从stem开始逐步下采样特征图从H/2到H/4一路降到H/32然后再通过上采样把分辨率升回来。这么做的好处是计算量可控低分辨率的特征图感受野大适合做分类这种全局任务。但是对像素级的任务来说空间细节被反复压缩再恢复信息损失是实打实的。比如关键点检测需要精确定位如果特征图缩到输入的1/32一个关键点的坐标偏差几个像素对应到原图上就是几十个像素精度很难保证。HRNet的思路反过来从第一个stage开始就保持一个高分辨率分支分辨率的下降不是串行的而是并行地不断添加低分辨率分支。网络主体部分高分辨率分支一直存在和低分辨率分支之间反复交换信息。高分辨率分支负责保留空间细节低分辨率分支负责提取语义信息通过反复的多尺度融合让两边都受益。这样输出特征图的空间分辨率始终维持在输入的1/4对于定位任务来说天然友好。1.2 并行多分支架构的核心逻辑HRNet的架构可以分几个阶段理解。第一个阶段就是高分辨率分支特征图尺寸是输入的1/4通道数和宽度由网络宽度参数决定。第二个阶段开始在保留第一个分支的基础上新增一个下采样到1/8的分支。第三个阶段再加一个1/16的分支第四个阶段再加1/32的分支。所以四个阶段结束时网络同时拥有从1/4到1/32四个不同分辨率的分支。每个阶段内部会做多次“multi-resolution parallel convolution”也就是多个分辨率的卷积并行执行并且在每个模块的输出处做多尺度信息交换。信息交换的方式是把每个分支的特征图下采样到其他分支的分辨率然后逐元素相加。低分辨率分支和高分辨率分支之间也会做上采样相加。这个交换过程让语义信息和空间信息可以互相补充这是HRNet区别于其他多尺度网络的本质。1.3 和U-Net、FPN这类多尺度网络有什么不同U-Net也有高分辨率特征但它是先下采样再上采样的对称结构高低分辨率特征只在跳跃连接时接触一次缺乏反复的信息融合。FPN是在ResNet基础上额外加了自顶向下的路径融合次数也有限。HRNet则是从始至终并行维护多分辨率特征每个stage都在做融合融合频率高得多。从实验结果看HRNet输出的特征质量更高尤其是对空间位置敏感的任务。另外HRNet没有像ResNet那样巨大的通道数变化它主要靠“多分支反复融合”来提升表征能力。计算量比同等精度的其他网络略高但换来的是更好的性能和更稳定的训练。实际用下来HRNet在小物体、遮挡场景下的稳定性明显更好。2. PyTorch环境准备与代码结构规划2.1 PyTorch安装与版本选择这部分对新手来说其实是最容易卡住的。PyTorch的安装强烈建议先去官网的Get Started页面选择对应的操作系统、包管理器、CUDA版本然后复制生成的命令来装。不要图省事用pip install torch一把梭因为这会把CPU版本装上GPU用不了。我自己常用的组合是PyTorch 2.x CUDA 11.8兼顾稳定性和新特性。如果用的是新显卡比如RTX 40系甚至50系建议选择CUDA 12.1及以上版本否则可能出现“no kernel image is available”的情况。装完之后用下面的命令验证是否能用GPUimport torch print(torch.__version__) print(torch.cuda.is_available())输出True就说明CUDA环境没问题。我在配置多台机器时发现最容易踩的坑是系统里装了好几个CUDA版本环境变量PATH和LD_LIBRARY_PATH指到了旧版本导致PyTorch编译出来的CUDA runtime和驱动不匹配。遇到这种情况先看nvidia-smi的驱动版本是否大于等于PyTorch要求的CUDA最低版本然后确认python -c import torch; print(torch.version.cuda)的版本和安装命令一致。2.2 项目代码整体结构实现HRNet时建议把代码拆成几个模块不要全部堆在一个文件里。我的目录结构是这样的hrnet_pytorch/ ├── models/ │ ├── __init__.py │ ├── hrnet.py # 主网络定义 │ ├── basic_blocks.py # 基础卷积block │ └── fuse_modules.py # 多分辨率融合模块 ├── config.py # 网络参数配置 ├── train.py # 训练入口 ├── dataset.py # 数据加载与预处理 └── utils.py # 工具函数这样拆的好处是以后想改网络结构或者换数据集只需要修改对应的模块不需要动其他部分。尤其是配置单独放一个文件调参数的时候特别方便不用翻网络定义。3. HRNet网络结构的核心代码实现3.1 基础模块Bottleneck与BasicBlockHRNet的stem部分通常用Bottleneck后面的并行分支用BasicBlock。Bottleneck和ResNet里的结构一致1x1卷积降维、3x3卷积、1x1卷积升维。BasicBlock则是两个3x3卷积中间带BN和ReLU。这里我直接把代码贴出来注释也写得比较完整import torch import torch.nn as nn def conv3x3(in_planes, out_planes, stride1): 3x3卷积same padding return nn.Conv2d(in_planes, out_planes, kernel_size3, stridestride, padding1, biasFalse) def conv1x1(in_planes, out_planes, stride1): 1x1卷积用于通道变换 return nn.Conv2d(in_planes, out_planes, kernel_size1, stridestride, biasFalse) class BasicBlock(nn.Module): expansion 1 def __init__(self, inplanes, planes, stride1, downsampleNone): super().__init__() self.conv1 conv3x3(inplanes, planes, stride) self.bn1 nn.BatchNorm2d(planes) self.relu nn.ReLU(inplaceTrue) self.conv2 conv3x3(planes, planes) self.bn2 nn.BatchNorm2d(planes) self.downsample downsample self.stride stride 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) if self.downsample is not None: identity self.downsample(x) out identity out self.relu(out) return out class Bottleneck(nn.Module): expansion 4 def __init__(self, inplanes, planes, stride1, downsampleNone): super().__init__() self.conv1 conv1x1(inplanes, planes) self.bn1 nn.BatchNorm2d(planes) self.conv2 conv3x3(planes, planes, stride) self.bn2 nn.BatchNorm2d(planes) self.conv3 conv1x1(planes, planes * self.expansion) self.bn3 nn.BatchNorm2d(planes * self.expansion) self.relu nn.ReLU(inplaceTrue) self.downsample downsample self.stride stride 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) out self.relu(out) out self.conv3(out) out self.bn3(out) if self.downsample is not None: identity self.downsample(x) out identity out self.relu(out) return out这里有一个细节Bottleneck的expansion是4所以输出通道是输入通道的四倍。stem部分用Bottleneck是为了初始特征提取更高效后面的基本模块用BasicBlock速度和显存占用都更可控。3.2 关键结构多分辨率并行卷积与融合HRNet的核心在于多个分支的并行卷积和融合。我先实现一个高通滤波器HighResolutionModule的类输入是多个分支的特征图列表输出也是同样数量的特征图列表但经过了一次并行卷积和融合。class HighResolutionModule(nn.Module): def __init__(self, num_branches, blocks, num_blocks, num_inchannels, num_channels, fuse_methodsum, multi_scale_outputTrue): super().__init__() self.num_branches num_branches self.num_inchannels num_inchannels self.fuse_method fuse_method self.multi_scale_output multi_scale_output self.branches self._make_branches( num_branches, blocks, num_blocks, num_inchannels, num_channels) self.fuse_layers self._make_fuse_layers() self.relu nn.ReLU(inplaceTrue) def _make_branches(self, num_branches, block, num_blocks, num_inchannels, num_channels): branches [] for i in range(num_branches): branches.append(self._make_one_branch( i, block, num_blocks, num_inchannels, num_channels)) return nn.ModuleList(branches) def _make_one_branch(self, branch_index, block, num_blocks, num_inchannels, num_channels): layers [] # 每个分支的第一个block可能需要下采样 downsample None if num_inchannels[branch_index] ! num_channels[branch_index]: downsample nn.Sequential( nn.Conv2d(num_inchannels[branch_index], num_channels[branch_index], kernel_size1, stride1, biasFalse), nn.BatchNorm2d(num_channels[branch_index])) layers.append(block(num_inchannels[branch_index], num_channels[branch_index], stride1, downsampledownsample)) num_inchannels[branch_index] num_channels[branch_index] for _ in range(1, num_blocks[branch_index]): layers.append(block(num_inchannels[branch_index], num_channels[branch_index])) return nn.Sequential(*layers)分支构建完成之后_make_fuse_layers负责生成融合所需的上下采样操作。融合规则是对于输出的第i个分支要把所有输入分支的特征图统一到第i个分支的分辨率。如果输入分支的分辨率比输出分支高就需要下采样比输出分支低就需要上采样。def _make_fuse_layers(self): num_branches self.num_branches num_inchannels self.num_inchannels fuse_layers [] for i in range(num_branches if self.multi_scale_output else 1): fuse_layer [] for j in range(num_branches): if j i: # 从低分辨率上采样到高分辨率 fuse_layer.append(nn.Sequential( nn.Conv2d(num_inchannels[j], num_inchannels[i], kernel_size1, stride1, biasFalse), nn.BatchNorm2d(num_inchannels[i]), nn.Upsample(scale_factor2 ** (j - i), modenearest))) elif j i: fuse_layer.append(None) else: # 从高分辨率下采样到低分辨率 conv3x3s [] for k in range(i - j): if k i - j - 1: num_outchannels_conv3x3 num_inchannels[i] conv3x3s.append(nn.Sequential( nn.Conv2d(num_inchannels[j], num_outchannels_conv3x3, kernel_size3, stride2, padding1, biasFalse), nn.BatchNorm2d(num_outchannels_conv3x3))) else: num_outchannels_conv3x3 num_inchannels[j] conv3x3s.append(nn.Sequential( nn.Conv2d(num_inchannels[j], num_outchannels_conv3x3, kernel_size3, stride2, padding1, biasFalse), nn.BatchNorm2d(num_outchannels_conv3x3), nn.ReLU(inplaceTrue))) fuse_layer.append(nn.Sequential(*conv3x3s)) fuse_layers.append(nn.ModuleList(fuse_layer)) return nn.ModuleList(fuse_layers)这里的下采样不是一次性把分辨率降到位而是逐步使用stride2的3x3卷积每降低一倍分辨率就做一次卷积。这样做比直接用一个stride4的卷积效果更好因为中间多了一层非线性而且参数量增加不多。融合的前向过程也比较关键需要特别小心张量形状的广播问题def forward(self, x): if self.num_branches 1: return [self.branches[0](x[0])] # 每个分支分别卷积 x [branch(branch_x) for branch, branch_x in zip(self.branches, x)] # 融合 x_fused [] for i in range(len(self.fuse_layers)): y x[0] if i 0 else self.fuse_layers[i][0](x[0]) for j in range(1, self.num_branches): if i j: y y x[j] else: # 保证channel匹配 y y self.fuse_layers[i][j](x[j]) x_fused.append(self.relu(y)) return x_fused值得注意的一点是在_make_fuse_layers中当从低分辨率上采样到高分辨率时用的是nearest插值而不是双线性插值。官方实现里用的就是nearest原因是上采样之后还有1x1卷积在学映射关系所以插值方式够用就行。如果换成bicubic反而增加了计算量精度提升也不明显。3.3 整体网络组装从stem到最终输出有了HighResolutionModule就可以搭完整的HRNet了。先定义stem然后按stage数量逐级增加分支。我这里以官方的HRNet-W32为例配置是stem输入是3通道先过一个3x3卷积(64通道)、BN、ReLU再接一个3x3卷积(64通道)、BN、ReLU输出分辨率是输入的1/4。stage11个Bottleneck通道64重复4次num_blocks[4]然后接一个HighResolutionModule。stage2新增1个低分辨率分支1/8两个分支都走BasicBlock每个分支重复4个blocknum_blocks[4, 4]。stage3三个分支分辨率分别是1/4、1/8、1/16每个分支重复4个blocknum_blocks[4, 4, 4]。stage4四个分支分辨率到1/32重复4个blocknum_blocks[4, 4, 4, 4]。在stage之间做一次transition也就是将现有的每个分支下采样一次作为下一个更低分辨率分支的初始化。下面给出简化的实现片段class HRNet(nn.Module): def __init__(self, config): super().__init__() self.config config # stem self.conv1 nn.Conv2d(3, 64, kernel_size3, stride2, padding1, biasFalse) self.bn1 nn.BatchNorm2d(64) self.conv2 nn.Conv2d(64, 64, kernel_size3, stride2, padding1, biasFalse) self.bn2 nn.BatchNorm2d(64) self.relu nn.ReLU(inplaceTrue) # 用于阶段间transition的层 num_channels config[num_channels] num_branches config[num_branches] # stage1 self.layer1 self._make_layer(Bottleneck, 64, 64, 4) # 后续stage self.stages nn.ModuleList() for i in range(1, len(num_branches)): self.stages.append(self._make_stage( num_branches[i], num_channels[i], num_channels[i-1], config[num_modules][i], config[num_blocks][i])) def _make_layer(self, block, inplanes, planes, blocks, stride1): downsample None if stride ! 1 or inplanes ! planes * block.expansion: downsample nn.Sequential( nn.Conv2d(inplanes, planes * block.expansion, kernel_size1, stridestride, biasFalse), nn.BatchNorm2d(planes * block.expansion)) layers [block(inplanes, planes, stride, downsample)] inplanes planes * block.expansion for _ in range(1, blocks): layers.append(block(inplanes, planes)) return nn.Sequential(*layers) def _make_stage(self, num_branches, num_channels, prev_channels, num_modules, num_blocks): modules [] # 第一个module需要处理transition的通道变化后续不需要 for i in range(num_modules): if i 0: in_channels prev_channels else: in_channels num_channels modules.append(HighResolutionModule( num_branches, BasicBlock, num_blocks, in_channels, num_channels)) return nn.Sequential(*modules)完整的forward中输入经过stem和前两个block之后得到初始特征然后依次进入各个stage。输出的x是一个列表包含当前所有分支的特征图。不同任务对输出的取法不一样分类任务通常取最后一个分支最低分辨率做全局平均池化接全连接层。姿态估计通常取第一个分支最高分辨率1/4作为热图预测的输入。语义分割可以取所有分支上采样到原图尺寸再concatenate。我在做姿态估计时直接取第一分支输出1/4分辨率接一个1x1卷积把通道数变成关键点数量然后双线性插值到原图尺寸作为热图。整个尾部很简单没有花里胡哨的结构。4. 训练配置与参数选择4.1 数据集准备与预处理用HRNet做姿态估计时常见的选择是COCO数据集。但完整COCO训练集有十几万张图对于只想快速验证网络是否跑通的情况太重了。我的建议是先拿一个小数据集比如MPII或者自造的简单关键点数据跑一版确认流程没问题再上COCO。预处理环节要重点注意几个点。第一图片输入尺寸建议统一到256x192这是HRNet官方在COCO上使用的分辨率。第二数据增强一般包括随机旋转正负30度、随机缩放0.75到1.25、随机平移、翻转。这些操作需要同步作用于关键点坐标和热图不是单纯处理图片就行。第三关键点坐标需要做归一化以图片尺寸为基准放到[-1, 1]或[0, 1]区间方便构造热图。生成热图的代码核心逻辑是把每个关键点映射到一个高斯核中心一般sigma取2def generate_heatmap(img_w, img_h, joints, num_joints, sigma2): heatmap np.zeros((num_joints, img_h // 4, img_w // 4), dtypenp.float32) for k in range(num_joints): x, y joints[k] if x 0 or y 0: continue # 下采样到1/4分辨率 x_hm int(x / 4) y_hm int(y / 4) # 限定范围 x_hm min(max(x_hm, 0), img_w // 4 - 1) y_hm min(max(y_hm, 0), img_h // 4 - 1) # 高斯核范围 ul int(x_hm - 3 * sigma), int(y_hm - 3 * sigma) br int(x_hm 3 * sigma), int(y_hm 3 * sigma) for i in range(ul[1], br[1] 1): for j in range(ul[0], br[0] 1): if 0 i img_h // 4 and 0 j img_w // 4: d (j - x_hm) ** 2 (i - y_hm) ** 2 heatmap[k, i, j] max( heatmap[k, i, j], np.exp(-d / (2 * sigma * sigma))) return heatmap这里的sigma2对应的是1/4分辨率的heatmap如果输入是256x192那么heatmap就是64x48。高斯核扩散半径设为3倍sigma既能保证峰值附近有足够的响应又不会让背景区域白白增加计算量。4.2 损失函数与优化器选择HRNet在姿态估计任务上最常用的损失函数是MSE均方误差直接比较预测热图和真实热图。MSE配合高斯热图非常自然因为热图本身就是一种概率分布的表达MSE相当于对高斯分布的每个位置做回归。很多人在这个环节会问为什么不用CrossEntropyLoss。实际对比下来MSE的训练更稳定收敛速度也更快。CrossEntropy更适合那种输出本身就是分类概率的情况而热图是回归目标用MSE是惯例。对于语义分割来说则一般用CrossEntropyLoss因为逐像素分类问题。优化器方面我用的是Adam初始学习率1e-3配合CosineAnnealing学习率调度。如果是COCO这种大数据集也可以使用SGDmomentummomentum0.9, weight_decay1e-4初始学习率设置为1e-2在epoch 90、120处做阶段性下降。SGD在大数据集上泛化能力稍好Adam在小数据集上起步更快。我自己的经验是先在小数据上用Adam跑通整个流程再切换SGD去跑大模型这样调试阶段省时间。4.3 训练过程中的关键参数监控训练过程中我重点看这几个指标loss下降曲线、当前epoch的PCK/OKS精度、每个分支特征的响应情况。光是loss在降还不够还要通过可视化确认关键点的预测是否合理地贴合人体结构。一个很有效的方法是每训练几个epoch保存一次预测热图和原图叠加后的可视化结果肉眼看比数值更直观。显存监控也很重要。HRNet-W32在batch size 64、输入256x192的情况下显存占用大概在10GB左右。如果显存只有8GB建议batch size降到32或者输入尺寸缩小到224x160。千万别硬撑OOM之后再临时改代码那样很浪费调试时间。5. 常见问题与排查技巧实录5.1 训练Loss不下降或下降极慢我自己最常遇到的第一个坑就是loss一开始就卡在某个值附近不动。排查路径一般按下面的顺序来检查数据归一化是否正确。输入图片如果是8-bit的uint8像素范围0-255需要除以255归一化到0-1否则初始loss会很大且收敛极慢。检查热图生成时的坐标对应关系。常见问题是在数据增强翻转后关键点坐标没有同步翻转导致热图位置错乱网络学不到有效信息。检查学习率是否设置合理。如果学习率太大loss曲线会出现剧烈震荡太小则下降幅度不明显。建议先用1e-3试跑100个iteration看loss有没有明显下降趋势。检查是否有出现NaN。如果用Adam且loss突然变成NaN大概率是学习率过高导致梯度爆炸或者BN层在batch size过小的情况下统计不稳定。解决方案是降低学习率并将batch size保持在8以上。5.2 训练效果可以但验证效果差这种情况多半是过拟合或者数据增强不够。小数据集上尤其常见。我的建议是增加随机的仿射变换强度、加入随机遮挡RandomOcclusion来模拟关键点被遮挡的场景。另外在验证阶段不要使用数据增强并且把模型切换到eval模式确保BN层使用running statistics而不是batch统计。还有一个容易被忽视的点是heatmap的下采样位置。关键点坐标如果是小数比如经过旋转后直接取整到整数像素位置会带来量化误差。正确做法是不要把关键点坐标除以4后直接取整而是在原图分辨率上生成高斯热图再用双线性插值或平均池化下采样到1/4。这样生成的热图峰值定位更准确网络学到的目标空间分布也更平滑。5.3 多卡训练时结果不一致HRNet在多卡训练时如果用了DistributedDataParallel需要注意每个进程的数据预处理必须完全一致尤其是随机数种子。最常见的问题是不同进程用了不同的随机增强参数导致模型看到的样本不一致。解决办法是在数据加载器里设置固定的seed或者在Dataset的__init__里为每个rank生成独立的随机数生成器。另外需要注意BN层在DDP模式下默认是统计全局的但如果你的总batch size比较小比如每张卡只有8张图那BN的统计可能不稳定。这种情况下可以考虑使用SyncBatchNorm把多个GPU的BN统计信息同步起来。HRNet的BN层数量不多换成SyncBN的额外开销可以接受。5.4 显存溢出问题显存溢出的解决方法按优先级排列第一减少batch size第二降低输入分辨率第三减少网络宽度比如把HRNet-W32换成HRNet-W18第四开启混合精度训练。PyTorch的自动混合精度AMP可以保持精度几乎不变的情况下节省约40%的显存。如果用了AMP之后仍然OOM那就需要仔细检查是否在backward过程中保持了很多中间变量。HRNet的并行分支数量多中间张量也多建议在用完特征之后及时释放比如x self.stages[i](x) # 不需要的分支及时del del x[2]不过一般情况下只要把batch size调到位是不需要做这种手动优化的。实在不行还可以用torch.utils.checkpoint对HighResolutionModule做梯度检查点用时间换空间但训练速度会慢20%到30%不推荐首选。6. 基于HRNet的扩展从姿态估计到其他任务6.1 语义分割中的HRNet用法HRNet做语义分割时通常保留四个分辨率的分支然后把所有分支的上采样到第一个分支的尺寸做1x1卷积对齐通道后concatenate再接一个分割头。这种做法的好处是既用了高分辨率的边界细节又用了低分辨率的语义信息。分割任务中HRNet和DeepLab系列相比优势在于不会产生过大的特征图计算量相对均衡。试验中HRNet-W48在Cityscapes数据集上能达到80%以上的mIoU和DeepLabV3差不多但参数量还少一些。不过要特别注意BN层在分割任务里对分辨率高度敏感如果训练和推理的输入尺寸不一致一定要先确认BN的统计量是否需要重新校准。6.2 目标检测中的HRNet替换Backbone在目标检测框架里替换backbone时HRNet的输出特征图一般是C2、C3、C4、C5对应1/4、1/8、1/16、1/32分辨率可以直接对接FPN。有实验表明把Faster R-CNN的ResNet-50换掉在COCO上AP能提升2到3个点但推理速度稍有下降主要原因是并行分支计算量高。在替换的时候要注意通道对齐。FPN期望输入通道数通常是256所以需要在HRNet的最后接一个1x1卷积做通道变换。另外HRNet的初始化方式建议使用ImageNet预训练权重不要从头训练否则收敛速度非常慢。官方在GitHub提供的预训练权重可以直接下载然后改改forward里的输出头就能用。6.3 轻量化版本与部署考虑HRNet的原始版本参数和计算量都不小如果要在移动端部署可以尝试HRNet-W18-Small或者把stem的stride改成1来提升分辨率保留程度。还有一种做法是把普通卷积改成深度可分离卷积能显著减少计算量。我自己试过在NPU设备上跑HRNet-W18-Small实时性勉强能达到但精度确实有所下滑。如果精度要求高建议直接换用更轻量的方案。部署阶段常遇到的问题是BN层和PyTorch模型的导出。导出ONNX时输入输出格式需要固定最好把HRNet封装成输入[N, 3, H, W]输出[N, num_joints, H/4, W/4]的单个输出。还有个坑是Upsample操作在ONNX导出时可能不支持某些插值模式建议先用nearest跑通导出流程再考虑优化。7. 我踩过的一些实战坑想单独拿出来说训练HRNet的过程中最大的一个教训是不要默认网络收敛很慢就直接改学习率。HRNet在骨干网络替换后因为BN层参数统计会重新适配前几个epoch的loss下降确实会比较慢但过了30个epoch后往往会突然加速。如果前20个epoch看到loss下降不明显就急着调大学习率很可能把网络搞发散。我当时在COCO子集上就吃过这个亏后来静下心跑了50个epoch才看到明显效果。另外给大家一个小技巧如果机器性能允许可以设置torch.backends.cudnn.benchmark True对固定输入尺寸的训练能提速20%到30%。但要注意如果输入尺寸不是固定的这个开关反而会导致程序变慢因为每次尺寸变化都要重新找最优卷积算法。HRNet如果固定输入尺寸比如256x192开启后效果很明显。最后再说一下可视化调试。训练初期我习惯把网络输出的特征图每隔一段时间打印出来看一眼确认不是一片模糊或大部分都为0。如果特征图大部分值都很小说明ReLU前的BN初始化可能有问题可以让BN的weight初始化为0.1而不是1这样初始激活值更活跃训练更容易启动。这个技巧在很多CNN任务里都实用。HRNet的实现并不复杂核心就是那几个并行分支和融合模块。只要理解了多分辨率并行和反复融合的思想代码本身还是很直观的。如果你在实现过程中卡住了优先检查张量形状和多分支数据的对应关系因为HRNet的索引比普通ResNet要多很多出错概率最大的就是分支不对齐。把基础版本跑通之后再往姿态估计、分割或者检测方向扩展都会顺很多。