libtorch工程化pytorch训练的模型实战

如果你想把pytorch训练好的模型快速部署到libtorch上，推荐通过把模型转化为torchscript脚本，然后在libtorch用C++接口加载这个模型。这个方式高效快速把训练模型和工程化快速对接起来。这个方式没有优化前行神经网络计算，但对于一般工程也是够用的,目前在centernet，3d卷积相关网络上成功过。

这种的原理大致是通过一次前向传播，把计算图保存下来，后面在python接口或者c++接口直接加载计算图，可以直接预测。

官方参考文档
https://pytorch.org/tutorials/advanced/cpp_export.html
https://pytorch.org/tutorials/advanced/torch_script_custom_ops.html
官方文档比较详细，但实战过程中对于复杂网络处理也有很多坑点。

一.下载pytorch 和libtorch
这种转化方式才出不久,一定要下载preview(预览版),标准1.1版本转化过程中有不少报错。
二.保存torchscript

import torch
import torchvision
model = torchvision.models.resnet18()
example = torch.rand(1, 3, 224, 224)
traced_script_module = torch.jit.trace(model, example)
traced_script_module.save("model.pt")

如代码所示，把模型和输入传进去就可以了,代码执行的时候会调用一次前向传播，然后生成网络图traced_script_module
　，然后保存为model.pt。
这里可能会遇到几个问题:
1.forward 的返回值支持有限，在转化centernet的时候,发现有个forward返回是python的list类型，这个转化过程会报出不支持此类型，我们需要把python代码手动改成tutple,然后在转化。通用map类型也要改成tuple。对于这种多输出网络结构，这种转化方式目前不支持这种可变类型的数据结构，tuple虽然也是多输出，但是个数树固定的。
2.多显卡转化过程中需要关闭
torch.nn.DataParallel(net.cuda(devices[0]), device_ids=[0])
如果代码开启了数据并行,转化可能失败，需要注释掉。
3.在转化pytorch扩展层的时候，可以把花里胡哨的写法干掉，比如在转化ceneternet可变形卷积的时候

 def forward(self, input, offset, mask):
        # assert 2 * self.deformable_groups * self.kernel_size[0] * self.kernel_size[1] == \
        #     offset.shape[1]
        # assert self.deformable_groups * self.kernel_size[0] * self.kernel_size[1] == \
        #     mask.shape[1]
        output = _backend.dcn_v2_forward(input, self.weight, self.bias,
                            offset, mask,
                            self.weight.shape[2], self.weight.shape[3],
                            self.stride[0], self.stride[1],
                            self.padding[0], self.padding[1],
                            self.dilation[0], self.dilation[1],
                            self.deformable_groups)
        return output
        # return dcn_v2_conv(input, offset, mask,
        #                    self.weight,
        #                    self.bias,
        #                    self.stride,
        #                    self.padding,
        #                    self.dilation,
        #                    self.deformable_groups)

把forward直接改成调用插件的函数,backward不用理会，工程化也只需要前向.
三.用python测试保存的torch scrip
这里有个简单的验证方法，代码里直把原来model替换一下,比如

module = torch.jit.load("model.pt")

再次启动成功如果和原来效果一样就成功了.
四.用c++　libtroch调用
参考官方文档，c++接口和python有些地方处理不一样，相关文档也比较少，可以看libtorch头文件慢慢学习。
一般上面pytorch script保存成功后面c++加载没什么问题.
对应插件层处理，可以把c++代码前向部分提取出来，注册成script函数,封成so，给ｃ＋＋调用。
如果你像center的可变形卷积用pytorch的cuda扩展写代码记得修改一个地方：

//extern THCState *state;
THCState *state = at::globalContext().lazyInitCUDA();

c++接口不会外部初始化pytorch cuda，这里设置一个懒加载，调用的时候初始化。