【他山之石】libtorch使用经验

最近花了点时间用pytorch的c++前端实现了faster rcnn的推理模型，在此分享一下libtorch的使用体验。

代码地址：

https://github.com/Mabinogiysk/ResNet50_libtorchfaster 

rcnn使用mmdetection2.2.1版本，未做任何修改；

c++的代码与mmdetection的inference代码一致，包括函数名和变量名，可以同步比对输入输出结果；

作者：阿卡丽

地址：https://www.zhihu.com/people/sky-3-42

01

先说说libtorch的使用体验和优缺点

1. 优点：

• libtorch前端完善程度超出我的预期，api非常友好，跟python很相似，模型转换的时候照着python代码一行一行翻译成c++就行。整个转换过程80%的时间都用来熟悉api，环境调试和代码编译，真正用来写代码的时间不多，一旦熟悉整个流程之后，模型转换就比较简单了；
• 简单的模型迁移非常方便，使用tracing或script保存后在c++中一行代码加载就能使用；
• 调试非常方便，直接跟python对各层的输出结果就行；

2. 缺点：

• 文档缺少示例，需要对c++有一定了解，否则可能看不懂；
• 教程匮乏，官方只有一个最简单的示例，对我帮助不大。期间遇到一个比较大的问题是不知道怎么把libtorch和.cu文件一起编译，这部分的CMakeLists我基本上是连蒙带猜写出来的，导致只能在cmake3.17上使用，此外libtorch1.6.0以下版本不能跟.cu文件一起编译；

性能方面，因为mmdetection中几个重要的性能瓶颈如nms和ROIAlign都已经用cuda实现，因此转到c++后性能提升很小，只快了5ms左右（30ms vs 35ms），不过cpu的占用率大幅下降。因此libtorch更多的使用场景应该是脱离python之后的部署，使用pytorch的extension和C++/CUDA足以解决大部分性能问题。

02

模型转转换

从pytorch转到libtorch有几种方式：

https://zhuanlan.zhihu.com/p/146453159

1. Tracing之后保存，适用于模型没有分支的情况；
2. Script编译，适用于模型有分支的情况，并且可以使用c++扩展；
3. 使用c++从零开始写模型，然后加载pytorch的权重；

方法3没有官方的api用来加载权重，需要自己写，非常麻烦，因此我在写了个Resnet50之后放弃了，最后使用类似方法2的办法完成了转换。我将faster rcnn分成了4个部分，纯network的部分使用Tracing来转换模型，其余部分使用c++，nms和ROIAlign使用cuda

Tracing部分

这个没什么好说的，官方的文档很详细。因为mmdetection的封装比较复杂，直接使用tracing不方便，我又单独写了一个faster rcnn的模型，用来加载mmdetection模型的权重，再用tracing保存为两个pt文件，分别对应上图的"backbone + fpn"和"CNN"部分。

NMS和ROIAlign

这两个直接使用了mmdetection的源码，难点在于编译和链接。直接使用nvcc编译ROIAlign代码会出错，需要在编译选项中加入"-D__CUDA_NO_HALF_OPERATORS__"阻止cuda使用自带的半精度计算

Debug

在python代码中直接将tensor保存为pt文件，然后使用如下代码将其转化为c++可以读取的格式：

class Container(torch.nn.Module):
    def __init__(self, my_values):
        super().__init__()
        for key in my_values:
            setattr(self, key, my_values[key])

my_values = {
    'img': torch.load('img.pt'),
    'proposal_list': torch.load('proposal_list.pt'),
    'cls_score': torch.load('cls_score.pt'),
    'bbox_pred': torch.load('bbox_pred.pt'),
    'det_bboxes': torch.load('det_bboxes.pt'),
    'det_labels': torch.load('det_labels.pt'),
}

# Save arbitrary values supported by TorchScript
# https://pytorch.org/docs/master/jit.html#supported-type
container = torch.jit.script(Container(my_values))
container.save("results.pt")

在c++使用如下代码加载results.pt文件并与c++结果比对：

torch::jit::script::Module results = torch::jit::load("results.pt");
Tensor std_img = results.attr("img").toTensor().to(device);
//auto feature_maps = results.attr("img").toTuple();
//auto cls_out = results.attr("img").toTensorVector();
std::cout << torch::sum(torch::abs(std_img - img)) << std::endl;

本文目的在于学术交流，并不代表本公众号赞同其观点或对其内容真实性负责，版权归原作者所有，如有侵权请告知删除。

直播预告

左划查看更多

历史文章推荐

分享、点赞、在看，给个三连击呗！