华为开源自研AI框架昇思MindSpore应用案例：人体关键点检测模型Lite-HRNet

如果你对MindSpore感兴趣，可以关注昇思MindSpore社区

一、环境准备

1.进入ModelArts官网

云平台帮助用户快速创建和部署模型，管理全周期AI工作流，选择下面的云平台以开始使用昇思MindSpore，获取安装命令，安装MindSpore2.0.0-alpha版本，可以在昇思教程中进入ModelArts官网

选择下方CodeLab立即体验

等待环境搭建完成

2.使用CodeLab体验Notebook实例

下载NoteBook样例代码 ，Mask-RCNN实现实例分割 ，.ipynb为样例代码

选择ModelArts Upload Files上传.ipynb文件

选择Kernel环境

切换至GPU环境，切换成第一个限时免费

进入昇思MindSpore官网，点击上方的安装

获取安装命令

回到Notebook中，在第一块代码前加入命令

conda update -n base -c defaults conda

安装MindSpore 2.0 GPU版本

conda install mindspore=2.0.0a0 -c mindspore -c conda-forge

安装mindvision

pip install mindvision

安装下载download

pip install download

人体关键点检测模型Lite-HRNet

人体关键点检测是计算机视觉的基本任务之一，在许多应用场景诸如自动驾驶、安防等有着重要的地位。可以发现，在这些应用场景下，深度学习模型可能需要部署在IoT设备上，这些设备算力较低，存储空间有限，无法支撑太大的模型，因此轻量但不失高性能的人体关键点检测级模型将极大降低模型部署难度。Lite-HRNet便提供了一轻量级神经网络骨干，通过接上不同的后续模型可以完成不同的任务，其中便包括人体关键点检测，在配置合理的情况下，Lite-HRNet可以以大型神经网络数十分之一的参数量及计算量达到相近的性能。

模型简介

Lite-HRNet由HRNet(High-Resolution Network)改进而来，HRNet的主要思路是在前向传播过程中通过维持不同分辨率的特征，使得最后生成的高阶特征既可以保留低分辨率高阶特征中的图像语义信息，也可以保留高分辨率高阶特征中的物体位置信息，进而提高在分辨率敏感的任务如语义分割、姿态检测中的表现。Lite-HRNet是HRNet的轻量化改进，改进了HRNet中的卷积模块，将HRNet中的参数量从28.5M降低至1.1M，计算量从7.1GFLOPS降低至0.2GFLOPS，但AP75仅下降了7%。

综上，Lite-HRNet具有计算量、参数量低，精度可观的优点，有利于部署在物联网低算力设备上服务于各个应用场景。

数据准备

本案例使用COCO2017数据集作为训练、验证数据集，请首先安装Mindspore Vision套件，并确保安装的Mindspore是GPU版本，随后请在https://cocodataset.org/ 上下载好2017 Train Images、2017 Val Images以及对应的标记2017 Train/Val Annotations，并解压至当前文件夹，文件夹结构下表所示

Lite-HRNet/
    ├── imgs
    ├── src
    ├── annotations
        ├──person_keypoints_train2017.json
        └──person_keypoints_train2017.json
    ├── train2017
    └── val2017

训练、测试原始图片如下所示，图片中可能包含多个人体，且包含的人体不一定包含COCO2017中定义的17个关键点，标注中有每个人体的边框、关键点信息，以便处理图像后供模型训练。

数据预处理

src/mindspore_coco.py中定义了供mindspore模型训练、测试的COCO数据集接口，在加载训练数据集时只需指定所用数据集文件夹位置、输入图像的尺寸、目标热力图的尺寸、以及手动设置对训练图像采用的变换即可

import mindspore as ms

import mindspore.dataset as dataset

import mindspore.dataset.vision.py_transforms as py_vision

import mindspore.nn as nn

from mindspore.dataset.transforms.py_transforms import Compose

from src.configs.dataset_config import COCOConfig

from src.dataset.mindspore_coco import COCODataset

cfg = COCOConfig(root="./", output_dir="outputs/", image_size=[192, 256], heatmap_size=[48, 64])

trans = Compose([py_vision.ToTensor(),

py_vision.Normalize(

mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])

train_ds = COCODataset(cfg, ".../", "train2017", True, transform=trans)

train_loader = dataset.GeneratorDataset(train_ds, ["data", "target", "weight"])

构建网络

Lite-HRNet网络骨干大体结构如下图所示：\

网络中存在不同的分辨率分支，网络主干上维持着较高分辨率、较少通道数的输出特征，网络分支上延展出较低分辨率、较多通道数的输出特征，且这些不同分辨率的特征之间通过上采样、下采样的卷积层进行交互、融合。Stage内的Cross Channel Weighting(CCW)则是网络实现轻量化的精髓，它将原HRNet中复杂度较高的1*1卷积以更低复杂度的Spatial Weighting等方法替代，从而实现降低网络参数、计算量的效果。CCW的结构如下图所示

首先Channel Split将输入特征通道分为两份，一份不作改动，另一份经过一次Cross Resolution Weighting、一次Spatial Weighting 后与前一份未作改动的特征进行沿通道方向的拼接，最后对拼接后的特征进行通道方向的重排以得到输出结果。Cross Resolution Weighting与Spatial Weighting的流程如下：

Cross Resolutional Weighting:

( X 1 , X 2 , . . . , X s ) − > A v g P o o l i n g − > C o n v . − > R e L U . − > C o n v . − > S i g m o i d − > U p s a m p l i n g − > ( W 1 , W 2 , . . . , W s ) (X_1,X_2,...,X_s) -> AvgPooling -> Conv. -> ReLU. -> Conv. -> Sigmoid -> Upsampling -> (W_1,W_2,...,W_s) (X1,X2,...,Xs)−>AvgPooling−>Conv.−>ReLU.−>Conv.−>Sigmoid−>Upsampling−>(W1,W2,...,Ws)

Y i = X i ⨀ W i Y_i=X_i \bigodot W_i Yi=Xi⨀Wi

Spatial Weighting:

X i − > A v g P o o l i n g − > F C . − > R e L U . − > F C . − > S i g m o i d − > W i X_i -> AvgPooling -> FC. -> ReLU. -> FC. -> Sigmoid -> W_i Xi−>AvgPooling−>FC.−>ReLU.−>FC.−>Sigmoid−>Wi

Y i = X i ⨀ W i Y_i=X_i \bigodot W_i Yi=Xi⨀Wi

上述公式中的 X i X_i Xi为第i个分辨率下的特征矩阵， W i W_i Wi为由 X i X_i Xi生成的、与 X i X_i Xi同形状的权重矩阵，由 X i X_i Xi生成 W i W_i Wi的过程分别对应图中的 H H H、 F F F函数，生成后的 W i W_i Wi与对应分辨率的 X i X_i Xi作哈马达积（即逐元素相乘）得到 Y i Y_i Yi对应图中的Channel Weighting部分，其中Cross Resolutional Weighting的各个 W i W_i Wi由 ( X 1 , X 2 , . . . , X s ) (X_1,X_2,...,X_s) (X1,X2,...,Xs)综合生成，所以会涉及上采样等操作。

而每个Stage后端存在Fusion Blocks以进行不同分辨率特征间交互融合，并产生更低分辨率的分支。下表列出了网络的具体结构参数。网络实际结构较复杂，具体细节请参考src/backbone.py中代码。

值得注意的是，除了骨干网络，作者在论文中同时也给出了所使用的检测头即SimpleBaseline，为了简洁起见，在本次的Lite-HRNet的Mindspore实现中，检测头（代码中包括IterativeHeads和LiteTopDownSimpleHeatMap）已集成至骨干网络之后，作为整体模型的一部分，直接调用模型即可得到热力图预测输出。

损失函数

此处使用损失函数为JointMSELoss，即关节点的均方差误差损失函数，其源码如下所示，总体流程即计算每个关节点预测热力图与实际热力图的均方差，其中target是根据关节点的人工标注坐标，通过二维高斯分布生成的热力图，target_weight用于指定参与计算的关节点，若某关节点对应target_weight取值为0，则表明该关节点在输入图像中未出现，不参与计算。

"""JointMSELoss"""

import mindspore.nn as nn

import mindspore.ops as ops

class JointsMSELoss(nn.Cell):

"""Joint MSELoss"""

def init (self, use_target_weight):

"""JointMSELoss"""

super(JointsMSELoss, self).init ()

self.criterion = nn.MSELoss(reduction='mean')

self.use_target_weight = use_target_weight

def construct(self, output, target, weight):
    """construct"""
    target = target
    target_weight = weight
    batch_size = output.shape[0]
    num_joints = output.shape[1]
    spliter = ops.Split(axis=1, output_num=num_joints)
    mul = ops.Mul()
    heatmaps_pred = spliter(output.reshape((batch_size, num_joints, -1)))
    heatmaps_gt = spliter(target.reshape((batch_size, num_joints, -1)))
    loss = 0

    for idx in range(num_joints):
        heatmap_pred = heatmaps_pred[idx].squeeze()
        heatmap_gt = heatmaps_gt[idx].squeeze()
        if self.use_target_weight:
            heatmap_pred = mul(heatmap_pred, target_weight[:, idx])
            heatmap_gt = mul(heatmap_gt, target_weight[:, idx])
            loss += 0.5 * self.criterion(
                heatmap_pred,
                heatmap_gt
            )
        else:
            loss += 0.5 * self.criterion(heatmap_pred, heatmap_gt)

    return loss/num_joints

模型实现与训练

在实现模型时，需指定模型内部结构，在src/net_configs中已指定原论文中10种结构配置，在训练样例种取Lite_18_coco作为模型结构，此处作为案例，仅设置epoch数量为1，在实际训练中可以设置为200，并且可以加入warmup。由于mindspore的训练接口默认数据集中每条数据只有两列（即训练数据和标签），所以这里需自定义Loss Cell。值得注意的是loss在训练前后变化并不会十分大，训练好的模型的loss为0.0004左右

class CustomWithLossCell(nn.Cell):

def __init__(self,
             net: nn.Cell,
             loss_fn: nn.Cell):
    super(CustomWithLossCell, self).__init__()
    self.net = net
    self._loss_fn = loss_fn

def construct(self, img, target, weight):
    """ build network """
    heatmap_pred = self.net(img)
    return self._loss_fn(heatmap_pred,
                         target,
                         weight)

from src.configs.net_configs import get_netconfig

from mindspore.train.callback import LossMonitor

from src.backbone import LiteHRNet

ext = get_netconfig("extra_lite_18_coco")

net = LiteHRNet(ext)

criterion = JointsMSELoss(use_target_weight=True)

train_loader = train_loader.batch(64)

optim = nn.Adam(net.trainable_params(), learning_rate=2e-3)

loss = JointsMSELoss(use_target_weight=True)

net_with_loss = CustomWithLossCell(net, loss)

model = ms.Model(network=net_with_loss, optimizer=optim)

epochs = 1

#Start Training

model.train(epochs, train_loader, callbacks=[LossMonitor(100)], dataset_sink_mode=False)

模型评估

模型评估过程中使用AP、AP50、AP75以及AR50、AR75作为评价指标，val2017作为评价数据集，pycocotool包中已实现根据评价函数，且src/mindspore_coco.py中的evaluate函数也实现了调用该评价函数的接口，只需提供预测关键点坐标等信息即可获得评价指标。此处载入Lite_18_coco的预训练模型进行评价。

from mindspore import load_checkpoint

from mindspore import load_param_into_net

from src.utils.utils import get_final_preds

import numpy as np

def evaluate_model(model, dataset, output_path):

"""Evaluate"""

num_samples = len(dataset)

all_preds = np.zeros(

(num_samples, 17, 3),

dtype=np.float32

)

all_boxes = np.zeros((num_samples, 6))
image_path = []

for i, data in enumerate(dataset):
    input_data, target, meta = data[0], data[1], data[3]
    input_data = ms.Tensor(input_data[0], ms.float32).reshape(1, 3, 256, 192)
    shit = model(input_data).asnumpy()
    target = target.reshape(shit.shape)
    c = meta['center'].reshape(1, 2)
    s = meta['scale'].reshape(1, 2)
    score = meta['score']
    preds, maxvals = get_final_preds(shit, c, s)
    all_preds[i:i + 1, :, 0:2] = preds[:, :, 0:2]
    all_preds[i:i + 1, :, 2:3] = maxvals
    # double check this all_boxes parts
    all_boxes[i:i + 1, 0:2] = c[:, 0:2]
    all_boxes[i:i + 1, 2:4] = s[:, 0:2]
    all_boxes[i:i + 1, 4] = np.prod(s*200, 1)
    all_boxes[i:i + 1, 5] = score
    image_path.append(meta['image'])

dataset.evaluate(0, all_preds, output_path, all_boxes, image_path)

net_dict = load_checkpoint("./ckpt/litehrnet_18_coco_256x192.ckpt")

load_param_into_net(net, net_dict)

eval_ds = COCODataset(cfg, "./", "val2017", False, transform=trans)

evaluate_model(net, eval_ds, "./result")

模型推理

1. Lite-HRNet是关键点检测模型，所以输入待推理图像应为包含单个人体的图像，作者在论文中提及在coco test 2017测试前已使用SimpleBaseline生成的目标检测Bounding Box处理图像，所以待推理图像应仅包含单个人体。
2. 网络的输入为(1,3,256,192)，所以在输入图像前应先将其变换成网络可处理的形式。
   import cv2
   from src.utils.utils import get_max_preds
   origin_img = cv2.imread("./imgs/man.jpg")
   origin_h, origin_w, _ = origin_img.shape
   scale_factor = [origin_w/192, origin_h/256]

resize to (112 112 3) and convert to tensor

img = cv2.resize(origin_img, (192, 256))

print(img.shape)

img = trans(img)

img = np.expand_dims(img, axis=0)

img = ms.Tensor(img)

print(img.shape)

Infer

heatmap_pred = net(img).asnumpy()

pred, _ = get_max_preds(heatmap_pred)

Postprocess

pred = pred.reshape(pred.shape[0], -1, 2)

print(pred[0])

pre_landmark = pred[0] * 4 * scale_factor

Draw points

for (x, y) in pre_landmark.astype(np.int32):

cv2.circle(origin_img, (x, y), 3, (255, 255, 255), -1)

Save image

cv2.imwrite("./imgs/man_infer.jpg", origin_img)

可以看到模型基本正确标注出了关键点的位置

算法基本流程

1. 获取原始数据
2. 从数据集的标注json文件中得到各个图像bbox以及关键点坐标信息
3. 根据bbox裁剪图像，并放缩至指定尺寸，如果是训练还可以作适当数据增强，生成指定尺寸的目标热力图
4. 指定尺寸的输入经过网络前向传播后得到预测的关键点热力图
5. 经过处理后取热力图中的最大值坐标作为关键点的预测坐标