Pytorch模型的C++ Libtorch推理：CenterNet

我的目标是在 Pytorch 中训练关键点检测的模型，在 Libtorch 中进行前向推理，实现 C++ 端的调用。在前向推理部分，目标检测和关键点检测的代码几乎一样，只是后处理部分不同。

在 Python 端利用 Pytorch 进行训练，然后使用 torch.jit 导出训练好的模型.pt

1. 训练模型

在 opts.py 中可以看出可以选择的网络架构有 res_18，dlav0_34，dla_34 和 hourglass 等，默认是dla_34，但它用到了dcn，转换成 Libtorch 比较复杂，所以我们在训练时使用 dlav0_34 作为网络架构。

2. 保存模型

我们写一个 save_model.py 用于保存训练好的模型：

from lib.models.model import create_model, load_model
import torch

if __name__ == '__main__':
    num_classes = 1   # 目标检测类别数
    head_conv = 256

   # 对照opt.py中的opt.heads，如果是目标检测任务则只有hm，wh和reg
    heads = {'hm': num_classes, 'wh': 2, 'hps': 10, 'reg': 2, 'hm_hp': 5, 'hp_offset': 2}
    load_model_path = 'model_best.pth'  # Pytorch中训练的pth模型
    save_script_pt = 'centernet.pt'  # 导出的pt模型，可以在C++中调用
    device = 0

    model = create_model('dlav0_34', heads, head_conv)
    model = load_model(model, load_model_path)
    model = model.to(device)
    model.eval()

    input_var = torch.zeros([1, 3, 512, 512], dtype=torch.float32).cuda()

    traced_script_module = torch.jit.trace(model, input_var)
    traced_script_module.save(save_script_pt)
    traced_script_module = torch.jit.load(save_script_pt)

但运行时会报错：RuntimeError: Only tensors, lists and tuples of tensors can be output from traced functions
原因是，我们看到 dlav0.py 文件中，DLASeg 类的 forward() 函数，返回的是字典 dict，我们需要把返回值修改成 tensor：

    def forward(self, x):
        x = self.base(x)
        x = self.dla_up(x[self.first_level:])
        ret = {}
        for head in self.heads:
            ret[head] = self.__getattr__(head)(x)
        return [ret]

需要修改成：

    def forward(self, x):
        x = self.base(x)
        x = self.dla_up(x[self.first_level:])
        temp_heads = []
        for head in self.heads:
            temp_head = self.__getattr__(head)(x)
            temp_heads.append(temp_head)

        re = torch.cat(temp_heads, 1)
        return re

CenterNet 的前向推理部分主要是数据的预处理，包括 rgb 值归一化和仿射变换等，
对于仿射变换：如果输入的图片大小是固定的，我们可以不在 C++ 里写 get_affine_transform 函数，而是直接用变换矩阵的值。比如我这里用的是 640x480 的图片，那么我们在 base_detector.py 中这里加一个 print()：

trans_input = get_affine_transform(c, s, 0, [inp_width, inp_height])
print(trans_input)

可以得到变换矩阵为：[[ 0.8 0 0 ], [ 0 0.8 64 ]]，在 C++ 中的 cv::warpAffine 函数中直接用这个变换矩阵值就可以。
C++ 中代码如下：

#include 
#include 
#include 
#include 
using namespace torch::indexing;

int main(int argc, char* argv[])
{

    //  调用关键点检测模型
	torch::jit::script::Module module = torch::jit::load("../centernet.pt");
	module.to(torch::kCUDA);
	
	int img_size = 512;
 
    std::vector inputs; 
 
    cv::Mat src, image, float_image;
    src = cv::imread("./image.jpg");

	//  仿射变换
    cv::Mat trans = (cv::Mat_(2, 3) << 0.8, 0, 0, 0, 0.8, 64);
	cv::warpAffine(src, image, trans, cv::Size(img_size, img_size), cv::INTER_LINEAR);

    //  归一化
    image.convertTo(float_image, CV_32F, 1.0 / 255);  

    auto img_tensor = torch::from_blob(float_image.data, {1, img_size, img_size, 3}, torch::kFloat32);  
    img_tensor = img_tensor.permute({0, 3, 1, 2});  

    // 参照自己数据集的均值和方差
    img_tensor[0][0] = img_tensor[0][0].sub_(0.485).div_(0.229);  
    img_tensor[0][1] = img_tensor[0][1].sub_(0.456).div_(0.224);
    img_tensor[0][2] = img_tensor[0][2].sub_(0.406).div_(0.225);


    auto img_var = torch::autograd::make_variable(img_tensor, false); 
    inputs.emplace_back(img_var.to(at::kCUDA)); 

    // 前向推理结果 result
    auto result = module.forward(inputs).toTensor();
    inputs.pop_back();
	
	return 0;
}