企业官网建设流程全解析

• 人工智能
• 计算机视觉
• 深度学习

导读

本文以 ViTPose 仓库（MMPose 生态）中configs/animal/2d_kpt_sview_rgb_img/topdown_heatmap/atrw/目录下的模型卡片 resnet_atrw.md 为骨架，系统讲解如何在 ATRW（东北虎野外重识别）数据集上使用 SimpleBaseline2D（ECCV'2018）方法配合 ResNet 骨干完成 2D 动物关键点估计。读完本文，你将掌握 ATRW 15 关键点任务的数据集定义、三档 ResNet 模型的完整配置语义（模型结构、优化器、学习率、数据增强、评估指标），并能够直接运行训练、测试与可视化推理命令。

一、背景：ATRW 数据集与 SimpleBaseline2D 方法

1.1 ATRW：野外东北虎重识别基准

ATRW（Amur Tiger Re-identification in the Wild）是 2020 年 ACM MM 发表的动物视觉基准数据集，论文标题为ATRW: A Benchmark for Amur Tiger Re-identification in the Wild，主要面向东北虎的个体重识别（Re-ID）任务，同时提供动物姿态（关键点）标注。其论文引用信息在模型卡片与数据集基类中均有记录：

• 模型卡片引用：resnet_atrw.md 中的li2020atrwBibTeX 条目；
• 数据集基类 atrw.py 中的paper_info字段；
• 数据加载类 animal_atrw_dataset.py 的 docstring 同样注明论文出处（arXiv:1906.05586）。

1.2 SimpleBaseline2D：顶向下热图方法的基线

方法侧，模型卡片引用了Simple Baselines for Human Pose Estimation and Tracking（ECCV'2018，作者 Bin Xiao 等），即常说的 SimpleBaseline2D。其核心思想非常简洁：使用 ImageNet 预训练的 ResNet 作为骨干提取特征，接若干反卷积（deconvolution）层恢复分辨率，最后用 1×1 卷积输出每个关键点的热图（heatmap），并以 MSE 损失监督。

在 ViTPose 仓库的代码体系中，该方法由TopDown检测器（top_down.py）+ResNet骨干（resnet.py）+TopdownHeatmapSimpleHead头部（topdown_heatmap_simple_head.py）组合实现。

二、配置结构总览：一份任务、三档模型

atrw目录下共有三份 ResNet 系列训练配置，均以256x256输入尺寸运行：

三份配置共享两个基类文件：

• default_runtime.py：运行时公共配置（checkpoint 保存间隔、日志、分布式参数等）；
• datasets/atrw.py：ATRW 数据集的dataset_info元信息（关键点定义、骨架、OKS sigmas）。

因此三份配置的差异几乎只集中在model.backbone.depth与model.pretrained两处，其余优化器、数据增强、评估设置完全一致——这正好体现了 SimpleBaseline 系列"换骨干、看效果"的消融式设计。

2.1 数据集元信息：15 个老虎关键点

atrw.py 中定义了 ATRW 数据集的 15 个关键点（keypoint_info），索引与名称如下：

0: left_ear 1: right_ear 2: nose 3: right_shoulder 4: right_front_paw 5: left_shoulder 6: left_front_paw 7: right_hip 8: right_knee 9: right_back_paw 10: left_hip 11: left_knee 12: left_back_paw 13: tail 14: center

• 每个关键点带有color（可视化颜色）与type（upper/lower，用于半身变换等数据增强）；
• swap字段描述左右对称关键点（如left_ear↔right_ear），用于水平翻转训练与翻转测试时关键点索引的重映射；
• skeleton_info定义了 14 条骨架连线（如left_ear–nose、tail–center），用于结果可视化；
• joint_weights全部为 1.0，表示各关键点在损失中的权重相同（use_different_joint_weights=False）；
• sigmas数组给出 15 个关键点的标注标准差，用于计算 OKS（Object Keypoint Similarity）指标，例如left_ear的 sigma 为 0.0277、nose为 0.0831。

在源码层，数据集类 animal_atrw_dataset.py 的 docstring 完整记录了同样的关键点索引表；基类 kpt_2d_sview_rgb_img_top_down_dataset.py 在初始化时会断言data_cfg['num_joints']与dataset_info.keypoint_num一致（L76），并将sigmas、skeleton、flip_pairs等元信息写入ann_info。

三、模型配置解析：TopDown + ResNet + 热图头

以 res50_atrw_256x256.py 为例，model字段结构如下：

channel_cfg = dict( num_output_channels=15, # 输出通道数 = 关键点数 dataset_joints=15, dataset_channel=[[0, 1, ..., 14]], # 数据集原始通道顺序 inference_channel=[0, 1, ..., 14]) # 推理时使用的通道顺序 model = dict( type='TopDown', pretrained='torchvision://resnet50', backbone=dict(type='ResNet', depth=50), keypoint_head=dict( type='TopdownHeatmapSimpleHead', in_channels=2048, # ResNet 最后一层输出通道 out_channels=channel_cfg['num_output_channels'], loss_keypoint=dict(type='JointsMSELoss', use_target_weight=True)), train_cfg=dict(), test_cfg=dict( flip_test=True, # 水平翻转测试融合 post_process='default', shift_heatmap=True, modulate_kernel=11))

要点逐一说明：

• 骨干（Backbone）：ResNet的depth在 {18, 34, 50, 101, 152} 中取值。从 resnet.py 的arch_settings（L478-L484）可以看到：50 层为Bottleneck结构(3, 4, 6, 3)，101 层为(3, 4, 23, 3)，152 层为(3, 8, 36, 3)。三档配置分别取 50/101/152，输出特征通道均为 2048，因此in_channels=2048在三份配置中保持不变。
• 预训练：torchvision://resnet50表示从 TorchVision 加载 ImageNet 预训练权重作为骨干初始化，这也是 SimpleBaseline 方法能快速收敛的关键。
• 头部（Head）：TopdownHeatmapSimpleHead是 SimpleBaseline 论文原版头部，由 3 层反卷积（每层 256 个滤波器、4×4 卷积核、stride 2，将特征图逐步上采样 8 倍）+ 1×1 卷积组成，最终输出 15 通道热图（源码见 topdown_heatmap_simple_head.py 的__init__默认参数num_deconv_layers=3, num_deconv_filters=(256, 256, 256), num_deconv_kernels=(4, 4, 4)）。
• 损失：JointsMSELoss配合use_target_weight=True，对每个关键点的热图预测与高斯目标之间计算加权 MSE；target_weight由关键点可见性决定。
• 测试策略：flip_test=True开启水平翻转测试（将原图与翻转图的热图结果按flip_pairs对齐后平均）；post_process='default'表示对热图做二阶导数插值求取亚像素坐标；shift_heatmap=True在翻转融合时校正热图偏移；modulate_kernel=11用于 max 反卷积调制。

3.1 数据配置（data_cfg）

data_cfg = dict( image_size=[256, 256], # 网络输入尺寸 heatmap_size=[64, 64], # 热图尺寸，为输入尺寸的 1/4 num_output_channels=15, num_joints=15, dataset_channel=[[0, ..., 14]], inference_channel=[0, ..., 14], soft_nms=False, # 评估时使用标准 OKS NMS nms_thr=1.0, oks_thr=0.9, # OKS 阈值 vis_thr=0.2, # 关键点可见性打分阈值 use_gt_bbox=True, # 训练/验证使用 GT 框（无需检测器） det_bbox_thr=0.0, bbox_file='', )

• image_size与heatmap_size的 4 倍关系对应上述 3 层 stride=2 反卷积的整体 8 倍上采样中前两步（骨干下采样 16 倍，头反卷积恢复 4 倍，热图相对输入为 1/4）；
• use_gt_bbox=True意味着该配置属于"以真实检测框为输入"的评估范式，即顶向下两阶段流程中检测阶段被 GT 框替代；
• 评估期的 NMS 参数（soft_nms、oks_thr、vis_thr）会传给数据集类的evaluate方法（详见第五节）。

四、训练配置：优化器、学习率与数据增强流水线

4.1 优化器与学习率调度

optimizer = dict(type='Adam', lr=5e-4) optimizer_config = dict(grad_clip=None) lr_config = dict( policy='step', warmup='linear', warmup_iters=500, warmup_ratio=0.001, step=[170, 200]) total_epochs = 210

• 使用Adam优化器，初始学习率5e-4，不做梯度裁剪；
• 学习率采用step 衰减：训练前 500 个 iteration 做线性 warmup（起始学习率为 0.001×lr），在第 170 与第 200 epoch分别衰减；
• 总训练轮数为210 epochs；
• evaluation = dict(interval=10, metric='mAP', save_best='AP')：每 10 个 epoch 在验证集上评估一次，指标为 mAP，并以 AP 为最佳模型保存依据（save_best='AP'）；
• log_config使用TextLoggerHook，每 1 个 iteration 记录一次日志。

4.2 数据加载与增强流水线

训练阶段（train_pipeline）依次执行：

train_pipeline = [ dict(type='LoadImageFromFile'), # 1. 读图 dict(type='TopDownRandomFlip', flip_prob=0.5), # 2. 随机水平翻转 dict(type='TopDownHalfBodyTransform', # 3. 半身变换 num_joints_half_body=8, prob_half_body=0.3), dict(type='TopDownGetRandomScaleRotation', # 4. 随机缩放旋转 rot_factor=40, scale_factor=0.5), dict(type='TopDownAffine'), # 5. 仿射变换 dict(type='ToTensor'), dict(type='NormalizeTensor', # 6. 归一化 mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), dict(type='TopDownGenerateTarget', sigma=2), # 7. 生成高斯热图目标 dict(type='Collect', keys=['img', 'target', 'target_weight'], meta_keys=['image_file', 'joints_3d', 'joints_3d_visible', 'center', 'scale', 'rotation', 'bbox_score', 'flip_pairs']), ]

细节说明：

• TopDownRandomFlip配合dataset_info中的flip_pairs同步翻转左右对称关键点；
• TopDownHalfBodyTransform以 0.3 的概率只保留上半身或下半身关键点（超过 8 个关键点的部分），增强对局部遮挡的鲁棒性；
• TopDownGetRandomScaleRotation使用 ±40° 旋转、±0.5 比例缩放；
• TopDownGenerateTarget以 sigma=2 生成高斯核热图目标；
• 归一化均值/标准差为 ImageNet 统计值，与 TorchVision 预训练权重匹配。

验证/测试阶段（val_pipeline与test_pipeline）则仅保留读图、仿射、归一化与 Collect，不进行随机增强。

数据加载器配置：

data_root = 'data/atrw' data = dict( samples_per_gpu=64, # 训练 batch size workers_per_gpu=2, val_dataloader=dict(samples_per_gpu=32), test_dataloader=dict(samples_per_gpu=32), train=dict(type='AnimalATRWDataset', ann_file=f'{data_root}/annotations/keypoint_train.json', img_prefix=f'{data_root}/images/train/', ...), val=dict(type='AnimalATRWDataset', ann_file=f'{data_root}/annotations/keypoint_val.json', img_prefix=f'{data_root}/images/val/', ...), test=dict(type='AnimalATRWDataset', ann_file=f'{data_root}/annotations/keypoint_val.json', img_prefix=f'{data_root}/images/val/', ...))

• 训练 batch size 为 64，验证/测试为 32；
• 数据集类统一为AnimalATRWDataset，训练用keypoint_train.json，验证/测试共用keypoint_val.json（仓库测试数据示例位于 tests/data/atrw/test_atrw.json）；
• 需要按data/atrw/annotations/...与data/atrw/images/...的目录结构自行准备 ATRW 数据集。

4.3 数据加载的源码支撑

AnimalATRWDataset继承自 kpt_2d_sview_rgb_img_top_down_dataset.py 中的Kpt2dSviewRgbImgTopDownDataset。其基类完成：

• 从 COCO 格式标注文件解析图像与关键点（COCO类来自 xtcocotools）；
• 将bbox（x, y, w, h）通过_xywh2cs（L132-L160）转换为(center, scale)表示，供仿射变换使用；
• 组合Compose(pipeline)完成数据流水线。

子类 animal_atrw_dataset.py 则负责：

• _get_db中assert self.use_gt_bbox，强制要求使用 GT 框；
• _load_coco_keypoint_annotation_kernel（L110-L176）逐图读取标注，过滤无效框（面积非正、越界等）与全零关键点，构造 db 条目（含center、scale、joints_3d、joints_3d_visible、bbox_score=1等）。

五、评估机制：mAP 指标与 OKS NMS 的实现

模型卡片的"Results on ATRW validation set"表是理解评估口径的关键。评估流程位于 animal_atrw_dataset.py 的evaluate方法（L178-L280），其allowed_metrics = ['mAP']，即 ATRW 任务只支持 COCO 风格的平均精度 mAP。流程如下：

1. 收集预测：将每个样本的preds（N×K×3，前两维为坐标、第三维为打分）、boxes、image_paths、bbox_ids汇总为 COCO 结果格式；
2. Rescoring（重打分）：对每个检测框，先统计超过vis_thr=0.2的关键点数量与平均关键点得分，得到kpt_score，最终score = kpt_score * box_score；
3. OKS NMS：use_nms=True且soft_nms=False时调用标准oks_nms（见 nms.py 的 L89-L128），按 OKS 重叠度去重，oks_thr=0.9为保留阈值；若开启soft_nms则改用soft_oks_nms（同文件 L155 起，基于高斯重打分）；
4. COCO 评估：将结果写入result_keypoints.json后调用COCOeval(self.coco, coco_det, 'keypoints', self.sigmas)（L329），sigmas来自 atrw.py 的 15 维数组，最终汇总输出 AP、AP .5、AP .75、AR 等 10 项统计量（L335-L338）。

值得注意的是，OKS 指标对每个关键点使用不同的 sigma（标注不确定性），因此在 ATRW 中耳朵、鼻子等关键点的 sigma 设置会直接影响 mAP 数值，这也是dataset_info必须精确配置的原因。

六、ATRW 验证集结果（来自模型卡片）

以下结果表完整继承自 resnet_atrw.md，输入尺寸均为 256×256，在 ATRW validation set 上评估：

观察结论（基于表格数据本身）：

• 三档骨干的 AP⁵⁰均为 0.973、AR⁵⁰均为 0.985，说明粗粒度定位已接近饱和；
• AP 与 AP⁷⁵三档模型非常接近（0.896~0.900），ResNet-50 略优，体现该任务上骨干深度增益有限，从部署性价比看ResNet-50 是首选基线。

对应的预训练权重与训练日志文件由 OpenMMLab 官方模型库托管，ckpt 文件名分别为res50_atrw_256x256-546c4594_20210414.pth、res101_atrw_256x256-da93f371_20210414.pth、res152_atrw_256x256-2bb8e162_20210414.pth，可通过 resnet_atrw.yml 中Models[].Weights字段定位下载入口（该 yml 同时记录了与上方表格一致的指标数值，供 model zoo 索引使用）。

七、训练、测试与推理实操

7.1 训练

使用 tools/train.py 即可启动训练，配置路径以仓库根目录为基准：

# 单卡训练 python tools/train.py \ configs/animal/2d_kpt_sview_rgb_img/topdown_heatmap/atrw/res50_atrw_256x256.py # 指定工作目录并覆盖配置（如修改 backbone 深度） python tools/train.py \ configs/animal/2d_kpt_sview_rgb_img/topdown_heatmap/atrw/res50_atrw_256x256.py \ --work-dir work_dirs/res50_atrw \ --cfg-options model.backbone.depth=101 model.pretrained=torchvision://resnet101

tools/train.py支持--resume-from断点续训、--no-validate关闭训练中验证、--seed/--deterministic固定随机性、--cfg-options覆盖配置（如--cfg-options model.backbone.depth=18），以及--launcher（none/pytorch/slurm/mpi）多机分布式训练。训练过程中每 10 个 epoch 会在验证集上计算 mAP 并保存最优模型。

7.2 测试

使用 tools/test.py 评估指定 checkpoint：

python tools/test.py \ configs/animal/2d_kpt_sview_rgb_img/topdown_heatmap/atrw/res50_atrw_256x256.py \ ${CHECKPOINT_FILE} \ --out result_keypoints.pkl

评估输出即上文第五节所述 COCO 风格 AP/AR 统计（包含AP .5、AP .75、AP (M)、AP (L)、AR等 10 项）。

7.3 可视化推理 Demo

由于 ATRW 配置使用 GT 框（use_gt_bbox=True），最直接的推理方式是提供包含 bbox 的 JSON 文件，调用 top_down_img_demo.py（用法模板见 2d_animal_demo.md 的"Using gt hand bounding boxes as input"一节）：

python demo/top_down_img_demo.py \ configs/animal/2d_kpt_sview_rgb_img/topdown_heatmap/atrw/res50_atrw_256x256.py \ ${CHECKPOINT_FILE} \ --img-root tests/data/atrw/ \ --json-file tests/data/atrw/test_atrw.json \ --out-img-root vis_results \ [--show] [--device cpu] \ [--kpt-thr 0.3]

• --img-root指向图片根目录，--json-file指向包含 bbox 的 COCO 格式 JSON；
• --device可指定cuda:0或cpu；
• --kpt-thr控制关键点可视化的置信度阈值，--radius、--thickness控制绘制样式；
• 输出可视化结果保存在--out-img-root目录下，骨架连线与关键点颜色由 atrw.py 中skeleton_info与keypoint_info.color决定。

八、测试用例：验证数据集与评估逻辑

仓库单元测试 test_animal_dataset.py 中test_animal_ATRW_dataset（L244-L302）直接复刻了本配置的channel_cfg与data_cfg，并使用tests/data/atrw/test_atrw.json构建AnimalATRWDataset，随后验证：

• dataset_name == 'atrw'、num_images == 2、__getitem__可正常返回样本；
• evaluate在 GT 输入下 mAP 应为 1.0（assert_almost_equal(infos['AP'], 1.0)）；
• 传入不支持的 metric（如PCK）会抛出KeyError，印证allowed_metrics=['mAP']的限制。

该测试从侧面确认了配置中dataset_info、data_cfg与数据集类的对接关系：基类 kpt_2d_sview_rgb_img_top_down_dataset.py 在dataset_info缺失时会告警，且必须从configs/_base_/datasets/atrw.py注入dataset_info（子类中亦有Config.fromfile('configs/_base_/datasets/atrw.py')的回退逻辑）。

九、小结

ATRW 老虎关键点估计是 ViTPose 仓库（MMPose 生态）中动物姿态任务的代表性配置之一。通过 resnet_atrw.md 这张模型卡片，可以完整追溯到：

• 数据侧：atrw.py 定义 15 关键点、14 条骨架与 OKS sigmas；animal_atrw_dataset.py 完成 COCO 格式加载与 mAP 评估；
• 模型侧：res50/res101/res152_atrw_256x256.py 给出 SimpleBaseline2D 顶向下热图方案的完整可复现配置；
• 指标侧：验证集 AP 约 0.90，且三档 ResNet 骨干表现接近，说明该任务基线已较为饱和，ResNet-50 可作为低成本首选。

对于后续研究者，从这份配置出发，可以自然扩展到其他动物数据集（如 horse10、macaque、fly 等同目录配置），或参考仓库中的 ViT 骨干（vit.py）将骨干替换为 Vision Transformer 以对比 CNN 基线，这也是本项目（ViTPose）的核心研究方向所在。

• 人工智能
• 计算机视觉
• 深度学习