目录
- 前言
- 1. alphapose导出
- 2. alphapose推理
- 3. 讨论
- 总结
前言
杜老师推出的 tensorRT从零起步高性能部署 课程,之前有看过一遍,但是没有做笔记,很多东西也忘了。这次重新撸一遍,顺便记记笔记。
本次课程学习 tensorRT 高级-alphapose模型导出、编译到推理(无封装)
课程大纲可看下面的思维导图
1. alphapose导出
这节课我们来学习 alphapose 姿态点估计
我们依旧从学习者的角度从零开始,拉取官方代码修改并正常导出 onnx,这个项目的复杂度较高,我们先看官方代码(下载于2022/3/27),使用的模型是 Fast Pose,如下图所示,由于 DCN 写插件比较麻烦,所以没有选择带 DCN的模型(目前 tensorRT_Pro 项目支持 DCN 算子)
先去执行 srcipts 脚本文件夹下的 demo_inference.py,看能否正常推理,执行如下:
可以看到出现 No module named ‘detector’ 错误,执行的 python 文件位于 scripts 文件夹下,而 detector 与 scripts 同级,有两种解决方式:
第一种就是 sys.path.insert() 插入 detector 的路径
第二种更加推荐就是 export PYTHONPATH=.
解决后再去执行,又报错了,如下所示:
出现 cython_bbox 错误,该模块主要用于 tracker 的,我们不需要 tracker 所以暂时屏蔽它,屏蔽后接着去执行,又报错了,如下所示:
出现 roi_align_cuda 错误,先屏蔽搁置后续如果需要再说,所以在 simple_transform.py 中屏蔽 21 和 22 行,再去执行,如下所示:
成功了,剩下的就是来提供参数,建议写一个脚本来输入参数,不要老是在终端去敲命令,infer.sh 内容如下所示:
#!/bin/bash
export PYTHONPATH=.
python scripts/demo_inference.py \
--cfg=configs/halpe_136/resnet/256x192_res50_lr1e-3_2x-regression.yaml \
--checkpoint=pretrained_models/multi_domain_fast50_regression_256x192.pth \
--sp \
--image=examples/demo/1.jpg \
--save_img
bash 执行下,又报错了,如下所示:
nms 有两套实现,我们直接屏蔽掉报错的那部分,直接使用手动实现,具体是 yolo_api.py 第 192 行,修改内容如下:
#if nms has to be done
if nms:
# if platform.system() != 'Windows':
# #We use faster rcnn implementation of nms (soft nms is optional)
# nms_op = getattr(nms_wrapper, 'nms')
# #nms_op input:(n,(x1,y1,x2,y2,c))
# #nms_op output: input[inds,:], inds
# _, inds = nms_op(image_pred_class[:,:5], nms_conf)
# image_pred_class = image_pred_class[inds]
# else:
# Perform non-maximum suppression
max_detections = []
while image_pred_class.size(0):
# Get detection with highest confidence and save as max detection
max_detections.append(image_pred_class[0].unsqueeze(0))
# Stop if we're at the last detection
if len(image_pred_class) == 1:
break
# Get the IOUs for all boxes with lower confidence
ious = bbox_iou(max_detections[-1], image_pred_class[1:], args)
# Remove detections with IoU >= NMS threshold
image_pred_class = image_pred_class[1:][ious < nms_conf]
image_pred_class = torch.cat(max_detections).data
然后屏蔽掉 yolo_api.py 的第 26 和 27 行,再去执行,又报错了,如下所示:
RoIAlign 没有定义,我们直接置为空,在 simple_transform.py 中的第 80 行修改为如下内容:
if platform.system() != 'Windows':
self.roi_align = None #RoIAlign(self._input_size, sample_num=-1)
if gpu_device is not None:
self.roi_align = self.roi_align.to(gpu_device)
再去执行,又报错了,如下所示:
接着去屏蔽,如下所示:
if platform.system() != 'Windows':
self.roi_align = None #RoIAlign(self._input_size, sample_num=-1)
# if gpu_device is not None:
# self.roi_align = self.roi_align.to(gpu_device)
那你可能会问为什么屏蔽?那凭感觉,通过理解罢了(还是得对整个流程熟悉呀😄)
再次执行,又报错了,如下所示:
上述问题是由于 pytorch 的模型名字不配对导致的,都是官方提供的但是报错了,说明官方没有做足够的 debug 才有这一堆破事,可以看到 checkpoint 的 shape 是 512,而 model 的 shape 是 1024,这说明刚才指定的配置文件和模型不匹配
我们来看下模型的定义,在 fastpost.py 中的第 44 行,错误提示 duc2.conv.weight 即 checkpoint 的 shape 是 512x256,而模型是 1024x256,还有 conv_out.weight 即 checkpoint 的 shape 是 136x128,而模型是 136x256,所以最终我们要把控制条件 slef.conv_dim 从 256 修改为 128,那这个变量是由 yaml 文件来控制的
因此我们需要从 256x192_res50_lr1e-3_2x-regression.yaml 文件中找到 CONV_DIM,将其从 256 修改为 128,再次执行,如下所示
成功了,说明我们修改的地方是正确的,模型推理的效果如下所示:
能正常推理了,接下来就是要把它正确的导出来了,在正式导出之前我们需要自己手动实现下推理过程,因此写一个 predict.py ,需要把它的整个推理过程像 Unet 一样抽出来,怎么去抽呢?官方推理都一堆 bug,这个事情就显得有些繁琐了
我们主要还是去参考 demo_inference.py 中的内容,根据各种分析最后得到的 predict.py 内容如下:
import yaml
from easydict import EasyDict as edict
from alphapose.models import builder
import torch
import numpy as np
import cv2
def update_config(config_file):
with open(config_file) as f:
config = edict(yaml.load(f, Loader=yaml.FullLoader))
return config
class MySPPE(torch.nn.Module):
def __init__(self):
super().__init__()
checkpoint = "pretrained_models/multi_domain_fast50_regression_256x192.pth"
cfg = update_config("configs/halpe_136/resnet/256x192_res50_lr1e-3_2x-regression.yaml")
self.pose_model = builder.build_sppe(cfg.MODEL, preset_cfg=cfg.DATA_PRESET)
self.pose_model.load_state_dict(torch.load(checkpoint, map_location="cpu"))
def forward(self, x):
hm = self.pose_model(x)
stride = int(256 / hm.size(2))
b, c, h, w = map(int, hm.size())
prob = hm.sigmoid()
confidence, _ = prob.view(-1, c, h * w).max(dim=2, keepdim=True)
prob = prob / prob.sum(dim=[2, 3], keepdim=True)
coordx = torch.arange(w, device=prob.device, dtype=torch.float32)
coordy = torch.arange(h, device=prob.device, dtype=torch.float32)
hmx = (prob.sum(dim=2) * coordx).sum(dim=2, keepdim=True) * stride
hmy = (prob.sum(dim=3) * coordy).sum(dim=2, keepdim=True) * stride
keypoint = torch.cat([hmx, hmy, confidence], dim=2)
return keypoint
model = MySPPE().eval()
x, y, w, h = 158, 104, 176, 693
image = cv2.imread("gril.jpg")[y:y+h, x:x+w]
image = image[..., ::-1]
image = cv2.resize(image, (256, 192))
image = ((image / 255.0) - [0.406, 0.457, 0.480]).astype(np.float32)
image = image.transpose(2, 0, 1)[None]
image = torch.from_numpy(image)
with torch.no_grad():
keypoint = model(image)
print(keypoint.shape)
#return torch.cat([hmx, hmy, confidence], dim=2)
dummy = torch.zeros(1, 3, 256, 192)
torch.onnx.export(
model, (dummy,), "fastpose.onnx", input_names=["image"], output_names=["predict"], opset_version=11,
dynamic_axes={
"image": {0:"batch"}, "predict": {0:"batch"}
}
)
print("Done")
杜老师通过分析把预处理和后处理给抽出来了,这要是自己分析不得疯,主要是代码封装得太深了,alphapose 的预处理部分在 simple_transform.py 文件中的 test_transform() 函数,输入一张原图和一个 box,进行相关预处理
整个预处理过程就是把 box 抠出来,然后移到中间,再减去均值就结束了。后处理部分是在 transforms.py 中的 heatmap_to_coord_simple_regress() 函数中实现的
值的注意的是我们在这里只是为了推理演示,去除了部分操作,我们直接拿一个已有的 box 塞到网络中去,省去了检测器的部分,同时拿到 box 后其实还是要做仿射变换的,这里为了方便直接使用的 resize,
模型推理结果是 136 维度的 heatmap,后处理就是是将 heatmap 变成回归值的过程,主要是得到我们的关键点坐标,这里把后处理部分也直接塞到 onnx 中,避免提高在 tensorRT 中的复杂度,
执行下 predict.py 如下所示:
导出的 onnx 如下图:
可以看到模型有很多多余的节点,都是 view 造成的,我们需要去除,在 SE_module.py 中 forward 部分修改,修改内容如下:
def forward(self, x):
b, c, _, _ = x.size()
# y = self.avg_pool(x).view(b, c)
y = self.avg_pool(x).view(-1, int(c))
# y = self.fc(y).view(b, c, 1, 1)
y = self.fc(y).view(-1, int(c), 1, 1)
return x * y
再导出下,onnx 如下图所示:
可以看到非常干净,是我们想要的效果
其实我们也可以直接拿 onnxsim 优化下,如下图所示:
import onnx
from onnxsim import simplify
onnx_model = onnx.load("fastpose.onnx")
model_simp, check = simplify(onnx_model)
onnx.save(model_simp, "fastpose.sim.onnx")
导出的 fastpose.sim.onnx 如下图所示:
清清爽爽,没有多余的节点,也非常 nice
2. alphapose推理
拿到我们想要的 onnx 后,接下来去 C++ 中执行下推理,直接 make run 运行下,如下所示:
来简单解读下代码,模型编译和之前没有任何区别,我们还是主要关注 inference,预处理部分通过 warpAffine 将图像缩放到 256x192,相比之前稍微做了下扩展,具体代码如下:
void get_preprocess_transform(const cv::Size& image_size, const cv::Rect& box, const cv::Size& net_size, float i2d[6], float d2i[6]){
cv::Rect box_ = box;
if(box_.width == 0 || box_.height == 0){
box_.width = image_size.width;
box_.height = image_size.height;
box_.x = 0;
box_.y = 0;
}
float rate = box_.width > 100 ? 0.1f : 0.15f;
float pad_width = box_.width * (1 + 2 * rate);
float pad_height = box_.height * (1 + 1 * rate);
float scale = min(net_size.width / pad_width, net_size.height / pad_height);
i2d[0] = scale; i2d[1] = 0; i2d[2] = -(box_.x - box_.width * 1 * rate + pad_width * 0.5) * scale + net_size.width * 0.5 + scale * 0.5 - 0.5;
i2d[3] = 0; i2d[4] = scale; i2d[5] = -(box_.y - box_.height * 1 * rate + pad_height * 0.5) * scale + net_size.height * 0.5 + scale * 0.5 - 0.5;
cv::Mat m2x3_i2d(2, 3, CV_32F, i2d);
cv::Mat m2x3_d2i(2, 3, CV_32F, d2i);
cv::invertAffineTransform(m2x3_i2d, m2x3_d2i);
}
这个 warpAffine 后的图像 input-image 如下所示,它其实是有一个扩大的过程,比我们平时的情况要复杂一点点
后处理由于我们是放在 onnx 的,因此直接获取的就是个关键点,根据置信度来进行过滤即可
检测效果如下图所示:
那这就是整个 alphapose 案例,有些地方看起来比较乱,还是需要自己多去实践,多去思考的,比如后处理就是这个算法的关键和核心,我们对官方代码进行解读后一定要自己实现一个版本,这样才能吸收消化从而变成我们自己的知识,还有一点,就是复杂的后处理放到 onnx 中可以解决很多问题
另外就是一个复杂的工程项目中要处理的问题太多了,但是我们要学会怎么化繁为简,这是我们要掌握的知识。
3. 讨论
姿态点估计算法可以分为自下而上和自上而下两种方法:(from chatGPT)
1. 自下而上方法:自下而上的姿态点估计算法是指先检测图像中所有可能的关键点,然后再通过关键点之间的关联关系来估计人体的姿态。这种方法通常从图像中检测出一系列的关键点,然后利用关键点之间的空间关系和约束关系来拟合出人体的姿态。
2. 自上而下方法:自上而下的姿态点估计算法是指先检测出人体的整体姿态或人体框,然后再在特定区域或人体框内估计关键点的位置。这种方法首先通过人体检测算法或目标检测算法找到人体的位置和姿态,然后在检测到的人体框内进行关键点的估计。
两种方法各有优势和适用场景:
- 自下而上方法的优势在于可以处理多人姿态估计问题,因为它能够检测图像中所有可能的关键点,然后通过关联关系对多人姿态进行建模。这种方法在密集场景中表现较好,但在处理复杂场景时可能存在误检或漏检问题。
- 自上而下方法的优势在于可以通过先验信息来辅助姿态估计,例如先进行人体检测或目标检测,然后再在检测到的人体框内进行关键点估计。这种方法通常比较高效,并且能够在复杂场景中保持稳定性,但可能不太适用于密集场景或多人姿态估计问题。
那很明显 alphapose 是自下而上的方法。
在 alphapose 中输入到网络中的是缩放到 256x192 尺寸的人体框,输出是一组热力图
在姿态点估计算法中,热力图(Heatmap)是一种用于表示关键点位置的图像。对于每个关键点,热力图是一个二维图像,其中每个像素的值表示该像素处是特定关键点的概率,热力图是如何转换成关键点坐标的呢?也就是后处理具体是如何做的呢?(这可能需要去仔细分析代码了😂)
那正常来说,整个 alphapose 的姿态点估计先通过检测器截取 box,再将截取得到的 box 送入到 alphapose 检测返回结果,如果人多的话,检测器截取到的每个 box 都要放到 alphapose 推理一遍,似乎有点耗时呀🤔
而且像多人密集场景,它则十分依赖检测器的能力,如果检测器的提取到的 box 不行,那后面的姿态点估计也就不准了,人少且比较分散效果应该不错
总结
这节课主要是学习姿态点估计网络 alphapose 的导出、编译到推理,这节课体现了一个非常重要的思想,那就是复杂的后处理放到 onnx 中去,这可以降低我们在 tensorRT 的复杂度。
同时这节课大部分时间都是在跟随杜老师不断解决各种各样的问题,我们实际工作中拿到一个工程项目文件也总是会遇到这样或者那样的问题,学习如何去解决问题才是我们要关注的,还是得多实践积累经验,能做到化繁为简,同时在理解完别人的代码后一定要自己实现一个版本,这样才能更好的去消化吸收变成我们自己的知识。