libtorch 上线之路-学习篇

最近在尝试将组里几个pytorch的模型使用libtorch 做nn inference，折腾了几个周，一路踩坑，终于把整个流程都打通了。因为libtorch还比较新，网上经验性的资料还比较少，现在写下来总结一下经验，分享一下。

0 环境

python 3.6
pytorch 1.3.1
g++ 5.3.1 # 要高于4.8，否则会报释放内存的错误
CentOS Linux release 7.2 (Final)
cuda 9.2

1 入门

参考pytorch官网的教程《LOADING A TORCHSCRIPT MODEL IN C++》.
网上大部分的文档都是对这篇教程的翻译。文档内容简单，一步一步的来就行了。

2 初窥门径

按照入门教程做完，只是知其然不知其所以然。接下来我们看看libtorch的原理。
libtorch c++和pytorch python这两门语言能互通的关键点是一个叫torchscript的东西。TorchScript可以理解为是pytorch的一个子集，能够支持部分python原生代码和部分pytorch接口，能够被libtorch c++读取和执行。
libtorch上线的原理就是将pytorch模型通过一定的方式转换为torchscript格式的模型，然后libtorch 读取torchscript格式的模型并执行。
所以把torchscript学透就差不多可以了。《TORCHSCRIPT》 官方教程以及《TORCHSCRIPT BUILTINS》torchscript支持的pytorch API，仔细多读几遍，差不多能够解决大部分的情况了。

3 打怪升级攒经验

将pytorch模型转为TorchScript格式的流程就是根据TorchScript的规则适应性的修改pytorch代码，并使用trace或者script模式将模型保存成TorchScript格式。

所以整个流程的核心其实就是能够将各种python代码修改成TorchScript能识别的格式。python这种具有各种包各种API的语言，大家经常会写出一些“随心所欲”的代码；再加上写python的人往往不关注效率二字；且libtorch的文档比较少。种种原因导致修改的难度增加不少。当然也看模型复杂度，比较简单的没有什么复杂处理的模型就没什么工作量了。

学完TorchScript的文档，只能说是能够应对一下常见的情况。想要做起来才能更得心应手，就需要你一个模型一个模型的打怪升级，总结经验。

在这里简单说一些我踩过的坑和总结的一些技巧。

3.1 不要使用除了原生python和pytorch以外的库

在nn inference的代码中（pytorch一般是forward函数和它调用的一些函数），不要使用除了原生python和pytorch以外的库，比如numpy什么的，TorchScript一律不支持……
python array虽然也支持，但也尽量少用，总之一句话**，尽量用pytorch的各种API**。

3.2 尽量使用traced模式将模型保存成TorchScript格式

TorchScript一共有两个模式将模型保存成TorchScript格式，traced和script。
script是直接将python代码固化起来。traced模式是记录python代码走过的路径，会进行一些优化操作（比如合并计算单元什么的），虽然现在基本没有啥优化操作。。但以后会有的。
所以在修改代码的时候，建议trace模式为主，script模式为副。

3.3 保证结果正确性

保证正确性是做nn inference的首先前提
在这里就是要保证TorchScript模型的结果跟训练代码的结果是一致的。
首先保证在python语言下结果是一致的，这个比较简单，demo 代码如下：

# 训练代码的结果
output = model(inputs)

# trace 模型的结果
trace_model = torch.jit.trace(model, (inputs))
trace_output = traced_model(inputs)

print(output.equal(trace_output))

然后还需要保证libtorch c++的结果跟python训练代码是一致的。
做法是将python中的输入和输出以一定格式写文件，然后在libtorchc++中读文件，再跟libtorch c++的结果进行对比。具体操作见 Load tensor from file in C++ fails

3.4 如何debug

这个是最烦心的。libtorch虽说是开源的，但是官方提供的库是编译好的，基本意味着gdb单步调试就很难了。只能用一些别的招。

1. TorchScript格式的模型是可以使用unzip进行解压的。解压后的目录大概如下：

project
│   constants.pkl
│   data.pkl 
│   version
└───code # TorchScript格式的代码
└───data # 权值
└───constants # 常量

code文件夹就是我们需要的，这里面可以看到TorchScript格式的python代码。当3.3对比结果不正确或者以后调优的时候，这个就比较有用了

2. pytorch github、论坛和教程。目前libtorch相关问题的资料基本都散落在这三个地方，耐心找吧……

3.5 一些libtorch api

libtorch c++ api文档的不完善程度简直令人发指。下面列一些我在头文件和网上找到的一些api用法

// tensor rand
torch::Tensor t = torch::rand({1, dim});
// Tensor赋值，注意类型
long sig_lens_arr[1] = {dim};
torch::Tensor t = torch::from_blob(sig_lens_arr, {1}, torch::kInt64);
// 获得 Tensor 的大小
cout << t.sizes()[0] << " " << t.sizes()[1] << endl;
// 打印Tensor
cout << t << endl;
// 访问Tensor内的值，注意类型
float*s_ptr = (float*)output_tensor.data_ptr();

结尾

把这篇博客里的东西都吃透了差不多就可以出新手村了，下一篇我们将进入性能篇，未完待续。
对了，招实习生啊！

【实习】【腾讯北京AILAB】招募AI异构加速实习生
简历直接给负责人，给简历保证迅速反馈。
基本条件: 熟悉c++，至少实习6个月
工作内容：
1. 使用c++复现框架训练的模型并进行CPU、GPU、ARM加速，达到上线的性能要求。
2. 调研各种inference框架并投入生产
加分项：
1、写过或者维护过深度学习框架代码； 
2、会CUDA 开发，能自己写kernel，会用cublas，cudnn等库； 
3、linux cpu c++编程能力，会写avx、会用mkl；
4、熟悉深度学习计算过程
5、学习能力强，实习时间长
联系方式: leowgyang@tencent.com