论文：《HigherHRNet: Scale-Aware Representation Learning for Bottom-Up Human Pose Estimation》

论文链接：https://arxiv.org/pdf/1908.10357.pdf

代码链接：https://github.com/HRNet/HigherHRNet-Human-Pose-Estimation

姿态估计前言知识：姿态估计-前言知识_error:404..的博客-CSDN博客

HigherHRNet详解之实验复现：HigherHRnet详解之实验复现_error:404..的博客-CSDN博客

HigherHRNet详解之论文详解：HigherHRnet详解之论文详解_error:404..的博客-CSDN博客

本文模型方面主要是在HRNet上进行了改进，加入了反卷积模块，那就先从模型上分析。

1.反卷积模块

    def _make_deconv_layers(self, cfg, input_channels):
        dim_tag = cfg.MODEL.NUM_JOINTS if cfg.MODEL.TAG_PER_JOINT else 1
        extra = cfg.MODEL.EXTRA
        deconv_cfg = extra.DECONV

        deconv_layers = []
        for i in range(deconv_cfg.NUM_DECONVS):
            if deconv_cfg.CAT_OUTPUT[i]:
                final_output_channels = cfg.MODEL.NUM_JOINTS + dim_tag \
                    if cfg.LOSS.WITH_AE_LOSS[i] else cfg.MODEL.NUM_JOINTS
                input_channels += final_output_channels
            output_channels = deconv_cfg.NUM_CHANNELS[i]
            deconv_kernel, padding, output_padding = \
                self._get_deconv_cfg(deconv_cfg.KERNEL_SIZE[i])

            layers = []
            layers.append(nn.Sequential(
                nn.ConvTranspose2d(
                    in_channels=input_channels,
                    out_channels=output_channels,
                    kernel_size=deconv_kernel,
                    stride=2,
                    padding=padding,
                    output_padding=output_padding,
                    bias=False),
                nn.BatchNorm2d(output_channels, momentum=BN_MOMENTUM),
                nn.ReLU(inplace=True)
            ))
            for _ in range(cfg.MODEL.EXTRA.DECONV.NUM_BASIC_BLOCKS):
                layers.append(nn.Sequential(
                    BasicBlock(output_channels, output_channels),
                ))
            deconv_layers.append(nn.Sequential(*layers))
            input_channels = output_channels

        return nn.ModuleList(deconv_layers)

 以及模型的初始参数;

DECONV.NUM_DCONVS = 2,
DECONV.NUM_CHANNELS = [32, 32],
DECONV.KERNEL_SIZE = [2, 2],
DECONV.NUM_BASIC_BLOCKS = 4
DECONV.CAT_OUTPUT = [True, True]

通过for构建两个DCONVS，两次都进行cat操作，拿一次构建来分析，

而后将 output_channels变成dconvs中的channel数，而后调用_get_deconv_cfg函数获取dconvs中的参数，具体是stride=2，kernel_size=2，padding=0，outpadding=0，

反卷积主要原理就是将一个像素块映射成多个像素块（stride>1时相当于插值操作，参考博客），就像下图那样，达到维度上升效果

 最终输入维度与输出维度关系：

Hout =  floor( Hin + 2*padding - kernel_size / stride) + 1

Hout就是想要得到的尺寸，Hin就是原来的尺寸，

Hout=floor（（Hin+2*0-2）/2）+1=2Hin

于是，也就是原文中说的通过反卷积模块可以将分辨率增加一倍。

然后再for构建4个BasicBlock残差模块。然后再input_channels = output_channels进行下一次构建。

2.train分析

valid跟train类似，以train为例进行分析。

2.1整体框架

dist_train.py

trian程序的入口，tools/dist_train.py 训练任务的执行文件，主函数。在dist_train.py中展示了HigherHRNet的整体流程。

dataset目录

target_generators.py: heatmap格式生成的函数，提供heatmap热图生成器。
transforms.py：数据增强的相关函数(沿用传统的翻转平移缩放等数据增强方法)

COCOKeypoints.py：被build.py直接调用。COCO数据集所用，是COCODataset.py的子类。COCOKeypoints.py会读取COCO的数据集，并进行数据增强，调用各种生成其制作所需的各类标签。最后一并输出img, target_list, mask_list, joints_list。

COCODataset.py：COCOKeypoints.py的父类，为COCOKeypoints.py提供基础的功能函数。

CrowdPoseKeypoints.py:CrowdPoseKeypoints.py会读取CrowdPose的数据集，格式和功能与COCOKeypoints.py一致。

CrowdPoseDataset.py:CrowdPoseKeypoints.py的父类，格式和功能与COCODataset.py一致。

build.py：封装data_loader，输入数据模块的一个功能主入口，调用target_generators.py，COCOKeypoints.py等文件的功能函数。最后直接输出模型训练所指定的数据格式。
 

pose_higher_hrnet.py

放着HigherHRNet的模型代码，主要是HRNet+2个dconvs（每个dconvs包含四个残差模块）

loss.py

存放着损失函数。
MultiLossFactory()是损失函数入口
其中主要包括热图损失(heatmaps_loss)和分组损失(ae_loss)。
分组方法采用了关联嵌入的分组法.

core/trainer.py

包含了整个训练的流程。通过dist_train.py文件调用函数do_train()来执行训练流程。其中包括了model的计算，然后输出的结果进行了loss计算，更新参数，最终输出log，保存模型等。

config目录

在训练和测试指令中都会用到cfg参数，实际上指向一个yaml文件。如果cfg下面代码已经设置了默认值，那么yaml里面的对应参数值会被覆盖，如果没有cfg下没有代码设置某个参数值，那么参数值就由yaml文件决定。

2.2 源码分析

2.2.1 dist_train.py

我们首先从tools/dist_train.py的主函数开始开始：
main函数中有各种设置。主要与流程相关的就是最后部分：是否采用分布式训练，若采用则调用多个main_worker函数；单卡训练则调用一次main_worker函数。main_worker可以看作时整个训练流程的入口。

def main():
    # 对输入参数进行解析
    args = parse_args()
    # 根据输入参数对默认的cfg配置进行更新
    update_config(cfg, args)
    
    cfg.defrost()
    cfg.RANK = args.rank
    cfg.freeze()
    # create logger
    logger, final_output_dir, tb_log_dir =create_logger(
        cfg, args.cfg, 'train'
    )

    logger.info(pprint.pformat(args))
    logger.info(cfg)

    # 判断输入参数中是否指定了GPU id
    if args.gpu is not None:
        warnings.warn('You have chosen a specific GPU. This will completely '
                      'disable data parallelism.')

    if args.dist_url == "env://" and args.world_size == -1:
        args.world_size = int(os.environ["WORLD_SIZE"])

    args.distributed = args.world_size > 1 or cfg.MULTIPROCESSING_DISTRIBUTED

    ngpus_per_node = torch.cuda.device_count()

	# 判断是否采用分布式训练
    if cfg.MULTIPROCESSING_DISTRIBUTED:
        # Since we have ngpus_per_node processes per node, the total world_size
        # needs to be adjusted accordingly
        args.world_size = ngpus_per_node * args.world_size
        # Use torch.multiprocessing.spawn to launch distributed processes: the
        # main_worker process function
        # 多采用分布式训练，会根据gpu数量启用相应个数的main_worker
        mp.spawn(
            main_worker,
            nprocs=ngpus_per_node,
            args=(ngpus_per_node, args, final_output_dir, tb_log_dir)
        )
    else:
        # Simply call main_worker function
        # 只需调用 main_worker 函数
        main_worker(
            ','.join([str(i) for i in cfg.GPUS]),
            ngpus_per_node,
            args,
            final_output_dir,
            tb_log_dir
        )

然后就进入main_worker

def main_worker(
        gpu, ngpus_per_node, args, final_output_dir, tb_log_dir
):
    # cudnn related setting
    # 使用GPU的一些相关设置
    cudnn.benchmark = cfg.CUDNN.BENCHMARK
    torch.backends.cudnn.deterministic = cfg.CUDNN.DETERMINISTIC
    torch.backends.cudnn.enabled = cfg.CUDNN.ENABLED

    # fp16 模式需要启用 cudnn 后端
    if cfg.FP16.ENABLED:
        assert torch.backends.cudnn.enabled, "fp16 mode requires cudnn backend to be enabled."

    # 警告：如果不使用 --fp16，static_loss_scale 将被忽略
    if cfg.FP16.STATIC_LOSS_SCALE != 1.0:
        if not cfg.FP16.ENABLED:
            print("Warning:  if --fp16 is not used, static_loss_scale will be ignored.")

    args.gpu = gpu

    if args.gpu is not None:
        print("Use GPU: {} for training".format(args.gpu))

    if args.distributed:
        if args.dist_url == "env://" and args.rank == -1:
            args.rank = int(os.environ["RANK"])
        if cfg.MULTIPROCESSING_DISTRIBUTED:
            # For multiprocessing distributed training, rank needs to be the
            # global rank among all the processes
            # 对于多进程分布式训练，rank需要是所有进程中的全局rank
            args.rank = args.rank * ngpus_per_node + gpu
        dist.init_process_group(
            backend=cfg.DIST_BACKEND,
            init_method=args.dist_url,
            world_size=args.world_size,
            rank=args.rank
        )
    # 根据输入参数对默认的cfg配置进行更新
    update_config(cfg, args)

    # setup logger
    logger, _ = setup_logger(final_output_dir, args.rank, 'train')

2.2.2 model文件

调用lib/models/pose_higher_hrnet.py的get_pose_net()函数，再初始化权重init_weights

2.2.3相关设置

包括设置日志，model的属性，以及多卡，单卡或cpu的选择。

    # copy model file
    if not cfg.MULTIPROCESSING_DISTRIBUTED or (
            cfg.MULTIPROCESSING_DISTRIBUTED
            and args.rank % ngpus_per_node == 0
    ):
        this_dir = os.path.dirname(__file__)
        shutil.copy2(
            os.path.join(this_dir, '../lib/models', cfg.MODEL.NAME + '.py'),
            final_output_dir
        )

    # 日志
    writer_dict = {
        'writer': SummaryWriter(log_dir=tb_log_dir),
        'train_global_steps': 0,
        'valid_global_steps': 0,
    }

    if not cfg.MULTIPROCESSING_DISTRIBUTED or (
            cfg.MULTIPROCESSING_DISTRIBUTED
            and args.rank % ngpus_per_node == 0
    ):
        dump_input = torch.rand(
            (1, 3, cfg.DATASET.INPUT_SIZE, cfg.DATASET.INPUT_SIZE)
        )

        # 输出日志
        writer_dict['writer'].add_graph(model, (dump_input, ))
        # logger.info(get_model_summary(model, dump_input, verbose=cfg.VERBOSE))

    # ???
    if cfg.FP16.ENABLED:
        model = network_to_half(model)

    if cfg.MODEL.SYNC_BN and not args.distributed:
        print('Warning: Sync BatchNorm is only supported in distributed training.')


    if args.distributed:
        if cfg.MODEL.SYNC_BN:
            model = nn.SyncBatchNorm.convert_sync_batchnorm(model)
        # For multiprocessing distributed, DistributedDataParallel constructor
        # should always set the single device scope, otherwise,
        # DistributedDataParallel will use all available devices.
        if args.gpu is not None:
            torch.cuda.set_device(args.gpu)
            model.cuda(args.gpu)
            # When using a single GPU per process and per
            # DistributedDataParallel, we need to divide the batch size
            # ourselves based on the total number of GPUs we have
            # 当每个进程和每个 DistributedDataParallel 使用单个 GPU 时，
            # 我们需要根据我们拥有的 GPU 总数自己划分批大小
            # args.workers = int(args.workers / ngpus_per_node)
            model = torch.nn.parallel.DistributedDataParallel(
                model, device_ids=[args.gpu]
            )
        else:
            model.cuda()
            # DistributedDataParallel will divide and allocate batch_size to all
            # available GPUs if device_ids are not set
            # 如果没有设置 device_ids，DistributedDataParallel 将划分并分配 batch_size 给所有可用的 GPU
            model = torch.nn.parallel.DistributedDataParallel(model)
    elif args.gpu is not None:
        # 单GPU工作
        torch.cuda.set_device(args.gpu)
        model = model.cuda(args.gpu)
    else:
        # 让模型支持CPU训练
        model = torch.nn.DataParallel(model).cuda()

2.2.4 损失函数

通过loss_factory = MultiLossFactory(cfg).cuda()调用core/loss.py的MultiLossFactory()函数,里面定义了热图损失函数和分组损失函数

# 热图损失
        self.heatmaps_loss = \
            nn.ModuleList(
                [
                    HeatmapLoss()
                    if with_heatmaps_loss else None
                    for with_heatmaps_loss in cfg.LOSS.WITH_HEATMAPS_LOSS
                ]
            )
        self.heatmaps_loss_factor = cfg.LOSS.HEATMAPS_LOSS_FACTOR
        # 分组损失
        self.ae_loss = \
            nn.ModuleList(
                [
                    AELoss(cfg.LOSS.AE_LOSS_TYPE) if with_ae_loss else None
                    for with_ae_loss in cfg.LOSS.WITH_AE_LOSS
                ]
            )

 2.2.5 训练

训练数据的生成模块，再main_worker只是一个简单的调用，但这之后有着这个源码最复杂的调用关系。

# Data loading code
    # 创建训练数据的迭代器
    train_loader = make_dataloader(
        cfg, is_train=True, distributed=args.distributed
    )
    logger.info(train_loader.dataset)

2.2.6  训练前的设置

if cfg.FP16.ENABLED:
    optimizer = FP16_Optimizer(
        optimizer,
        static_loss_scale=cfg.FP16.STATIC_LOSS_SCALE,
        dynamic_loss_scale=cfg.FP16.DYNAMIC_LOSS_SCALE
    )

begin_epoch = cfg.TRAIN.BEGIN_EPOCH     # 0
checkpoint_file = os.path.join(
    final_output_dir, 'checkpoint.pth.tar')  # 加载中间节点的模型权重信息（接续训练）
if cfg.AUTO_RESUME and os.path.exists(checkpoint_file):
    logger.info("=> loading checkpoint '{}'".format(checkpoint_file))
    checkpoint = torch.load(checkpoint_file)    # 载入节点权重信息
    begin_epoch = checkpoint['epoch']   # 开始的批次数
    best_perf = checkpoint['perf']  # 最优的结果？
    last_epoch = checkpoint['epoch']    #
    # 加载模型权重
    # state_dict（） 是一个Python字典，将每一层映射成它的参数张量。
    # 注意只有带有可学习参数的层（卷积层、全连接层等），以及注册的缓存（batchnorm的运行平均值）在state_dict 中才有记录
    model.load_state_dict(checkpoint['state_dict'])
    # 加载权重
    optimizer.load_state_dict(checkpoint['optimizer'])
    logger.info("=> loaded checkpoint '{}' (epoch {})".format(
        checkpoint_file, checkpoint['epoch']))

if cfg.FP16.ENABLED:
    # pytorch动态调整学习率
    lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
        optimizer.optimizer, cfg.TRAIN.LR_STEP, cfg.TRAIN.LR_FACTOR,
        last_epoch=last_epoch
    )
else:
    # pytorch动态调整学习率
    lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
        optimizer, cfg.TRAIN.LR_STEP, cfg.TRAIN.LR_FACTOR,
        last_epoch=last_epoch
    )

2.2.7 训练部分

# 开始训练
    # Epoch：一个Epoch就是将所有训练样本训练一次的过程。
    # Batch（批 / 一批样本）：将整个训练样本分成若干个Batch。
    # Batch_Size（批大小）： 每批样本的大小。
    # Iteration（一次迭代）：训练一个Batch就是一次Iteration。
    for epoch in range(begin_epoch, cfg.TRAIN.END_EPOCH):
        count+=1
        print("length:",len(range(begin_epoch, cfg.TRAIN.END_EPOCH)))
        print("count:",count)

        # train one epoch
        do_train(cfg, model, train_loader, loss_factory, optimizer, epoch,
                 final_output_dir, tb_log_dir, writer_dict, fp16=cfg.FP16.ENABLED)

        # In PyTorch 1.1.0 and later, you should call `lr_scheduler.step()` after `optimizer.step()`.
        # 是用来更新优化器的学习率的，一般是按照epoch为单位进行更换，
        # 即多少个epoch后更换一次学习率，因而scheduler.step()放在epoch这个大循环下。
        lr_scheduler.step()

        perf_indicator = epoch
        # 更新最优模型
        if perf_indicator >= best_perf:
            best_perf = perf_indicator
            best_model = True
        else:
            best_model = False

        if not cfg.MULTIPROCESSING_DISTRIBUTED or (
                cfg.MULTIPROCESSING_DISTRIBUTED
                and args.rank == 0
        ):

            logger.info('=> saving checkpoint to {}'.format(final_output_dir))

            # checkpoint检查点：不仅保存模型的参数，优化器参数，还有loss，epoch等（相当于一个保存模型的文件夹）
            # 在保存用于推理或者继续训练的常规检查点的时候，除了模型的state_dict之外，还必须保存其他参数。
            save_checkpoint({
                'epoch': epoch + 1,
                'model': cfg.MODEL.NAME,
                'state_dict': model.state_dict(),
                'best_state_dict': model.module.state_dict(),
                'perf': perf_indicator,
                'optimizer': optimizer.state_dict(),
            }, best_model, final_output_dir)

    final_model_state_file = os.path.join(
        final_output_dir, 'final_state{}.pth.tar'.format(gpu)
    )

    logger.info('saving final model state to {}'.format(
        final_model_state_file))

    # 保存整个模型
    torch.save(model.module.state_dict(), final_model_state_file)
    writer_dict['writer'].close()

里面有些需要注意的

执行下面代码，会触发之前实例化的dataset的一个函数__getitem__()

for i, (images, heatmaps, masks, joints) in enumerate(data_loader):

在CocoKeypoints中可以找到，在enumerate枚举时会触发。
调用CocoKeypoints的__getitem__()同时也会调用父类CocoDataset的__getitem__()

最终CocoKeypoint的__getitem__()会返回4个返回值。img, target_list, mask_list, joints_list就是对应着trainer.py中的(images, heatmaps, masks, joints)

最终会计算损失，

loss_factory实际上就是之前提到的MultiLossFactory函数。

实际上在执行MultiLossFactory的forward函数。

 最后再对损失函数进行一个综合计算，得出总损失值：