作者火星少女@知乎
来源丨https://zhuanlan.zhihu.com/p/146453159
极市平台编辑
最近，由于工作需要，我们必须把它放在一边pytorch的模型部署到c 在平台上，基本流程主要参照官网的教学示例，在此期间发现了许多坑，特此记录。
1.模型转换
libtorch不依赖于python，python训练模型需要转换为script model才能由libtorch加载，推理。这一步官网提供了两种方法:
这种方法操作相对简单，只需要一组输入模型，再次推理网络，然后从torch.ji.trace记录路径上的信息并保存它例如：
import torch import torchvision # An instance of your model. model = torchvision.models.resnet18() # An example input you would normally provide to your model's forward() method. example = torch.rand(1, 3, 224, 224) # Use torch.jit.trace to generate a torch.jit.ScriptModule via tracing. traced_script_module = torch.jit.trace(model, example)缺点是如果模型中有控制流，例如if-else句子，一组输入只能通过一个分支，在这种情况下不能完全记录模型信息。
直接在Torch在脚本中编写模型并相应注释模型torch.jit.script将编译模块转换为编译模块ScriptModule。示例如下：
class MyModule(torch.nn.Module): def __init__(self, N, M): super(MyModule, self).__init__() self.weight = torch.nn.Parameter(torch.rand(N, M)) def forward(self, input): if input.sum() > 0: output = self.weight.mv(input) else: output = self.weight input return output my_module = MyModule(10,20) sm = torch.jit.script(my_module)forward默认编译方法，forward调用的方法也会按照调用的顺序编译
若要编译一个forward以外且未被forward可以添加调用方法
如果不需要编译方法，可以使用
（https://pytorch.org/docs/master/generated/torch.jit.ignore.html#torch.jit.ignore）
或者（https://pytorch.org/docs/master/generated/torch.jit.unused.html#torch.jit.unused）
# Same behavior as pre-PyTorch 1.2 @torch.jit.script def some_fn(): return 2 # Marks a function as ignored, if nothing # ever calls it then this has no effect @torch.jit.ignore def some_fn2(): return 2 # As with ignore, if nothing calls it then it has no effect. # If it is called in script it is replaced with an exception. @torch.jit.unused def some_fn3(): import pdb; pdb.set_trace() return 4 # Doesn't do anything, this function is already # the main entry point @torch.jit.export def some_fn4(): return 2这一步遇到很多坑的主要原因可以分为两点
TorchScript支持的操作是python的子集，大部分torch使用的操作可以找到相应的实现，但也有一些尴尬的不支持操作，详细列表可见https://pytorch.org/docs/master/jit_unsupported.html#jit-unsupported，以下是我自己遇到的一些操作:
1）参数/返回值不支持可变个数，例如
def __init__(self, **kwargs):或者
if output_flag == 0: return reshape_logits else: loss = self.loss(reshape_logits, term_mask, labels_id) return reshape_logits, loss2)各种iteration操作
eg1.
layers = [int(a) for a in layers]报错torch.jit.frontend.UnsupportedNodeError: ListComp aren’t supported
可以改成：
for k in range(len(layers)): layers[k] = int(layers[k])eg2.
seq_iter = enumerate(scores) try: _, inivalues = seq_iter.__next__() except: _, inivalues = seq_iter.next()eg3.
line = next(infile)3)不支持的句子
eg1. 不支持continue
torch.jit.frontend.UnsupportedNodeError: continue statements aren’t supported
eg2. 不支持try-catch
torch.jit.frontend.UnsupportedNodeError: try blocks aren’t supported
eg3. 不支持with语句
4）其他常见op/module
eg1. torch.autograd.Variable
解决：使用torch.ones/torch.randn等初始化+.float()/.long()等指定数据类型。
eg2. torch.Tensor/torch.LongTensor etc.
解决：同上
eg3. requires_grad参数只在torch.tensor中支持，torch.ones/torch.zeros等不可用
eg4. tensor.numpy()
eg5. tensor.bool()
解决：tensor.bool()用tensor>0代替
eg6. self.seg_emb(seg_fea_ids).to(embeds.device)
解决：需要转gpu的地方显示调用.cuda()
总之一句话：除了原生python和pytorch以外的库，比如numpy什么的能不用就不用，尽量用pytorch的各种API。
2. 指定数据类型
1）属性，大部分的成员数据类型可以根据值来推断，空的列表/字典则需要预先指定
from typing import Dict
class MyModule(torch.nn.Module):
my_dict: Dict[str, int]
def __init__(self):
super(MyModule, self).__init__()
# This type cannot be inferred and must be specified
self.my_dict = {}
# The attribute type here is inferred to be `int`
self.my_int = 20
def forward(self):
pass
m = torch.jit.script(MyModule())2）常量，使用Final关键字
try:
from typing_extensions import Final
except:
# If you don't have `typing_extensions` installed, you can use a
# polyfill from `torch.jit`.
from torch.jit import Final
class MyModule(torch.nn.Module):
my_constant: Final[int]
def __init__(self):
super(MyModule, self).__init__()
self.my_constant = 2
def forward(self):
pass
m = torch.jit.script(MyModule())3）变量。默认是tensor类型且不可变，所以非tensor类型必须要指明
def forward(self, batch_size:int, seq_len:int, use_cuda:bool):方法三：Tracing and Scriptin混合
一种是在trace模型中调用script，适合模型中只有一小部分需要用到控制流的情况，使用实例如下：
import torch
@torch.jit.script
def foo(x, y):
if x.max() > y.max():
r = x
else:
r = y
return r
def bar(x, y, z):
return foo(x, y) + z
traced_bar = torch.jit.trace(bar, (torch.rand(3), torch.rand(3), torch.rand(3)))另一种情况是在script module中用tracing生成子模块，对于一些存在script module不支持的python feature的layer，就可以把相关layer封装起来，用trace记录相关layer流，其他layer不用修改。使用示例如下：
import torch
import torchvision
class MyScriptModule(torch.nn.Module):
def __init__(self):
super(MyScriptModule, self).__init__()
self.means = torch.nn.Parameter(torch.tensor([103.939, 116.779, 123.68])
.resize_(1, 3, 1, 1))
self.resnet = torch.jit.trace(torchvision.models.resnet18(),
torch.rand(1, 3, 224, 224))
def forward(self, input):
return self.resnet(input - self.means)
my_script_module = torch.jit.script(MyScriptModule())2.保存序列化模型
如果上一步的坑都踩完，那么模型保存就非常简单了，只需要调用save并传递一个文件名即可，需要注意的是如果想要在gpu上训练模型，在cpu上做inference，一定要在模型save之前转化，再就是记得调用model.eval(),形如
gpu_model.eval()
cpu_model = gpu_model.cpu()
sample_input_cpu = sample_input_gpu.cpu()
traced_cpu = torch.jit.trace(traced_cpu, sample_input_cpu)
torch.jit.save(traced_cpu, "cpu.pth")
traced_gpu = torch.jit.trace(traced_gpu, sample_input_gpu)
torch.jit.save(traced_gpu, "gpu.pth")3.C++ load训练好的模型
要在C ++中加载序列化的PyTorch模型，必须依赖于PyTorch C ++ API（也称为LibTorch）。libtorch的安装非常简单，只需要在pytorch官网（https://pytorch.org/）下载对应版本，解压即可。会得到一个结构如下的文件夹。
libtorch/
bin/
include/
lib/
share/然后就可以构建应用程序了，一个简单的示例目录结构如下：
example-app/
CMakeLists.txt
example-app.cppexample-app.cpp和CMakeLists.txt的示例代码分别如下：
#include <torch/script.h> // One-stop header.
#include <iostream>#include <memory>
int main(int argc, const char* argv[]) {
if (argc != 2) {
std::cerr << "usage: example-app <path-to-exported-script-module>\n";
return -1;
}
torch::jit::script::Module module;
try {
// Deserialize the ScriptModule from a file using torch::jit::load().
module = torch::jit::load(argv[1]);
}
catch (const c10::Error& e) {
std::cerr << "error loading the model\n";
return -1;
}
std::cout << "ok\n";
}cmake_minimum_required(VERSION 3.0 FATAL_ERROR)
project(custom_ops)
find_package(Torch REQUIRED)
add_executable(example-app example-app.cpp)
target_link_libraries(example-app "${TORCH_LIBRARIES}")
set_property(TARGET example-app PROPERTY CXX_STANDARD 14)至此，就可以运行以下命令从example-app/文件夹中构建应用程序啦：
mkdir build
cd build
cmake -DCMAKE_PREFIX_PATH=/path/to/libtorch ..
cmake --build . --config Release其中/path/to/libtorch是之前下载后的libtorch文件夹所在的路径。这一步如果顺利能够看到编译完成100%的提示，下一步运行编译生成的可执行文件，会看到“ok”的输出，可喜可贺！
4. 执行Script Module
终于到最后一步啦！下面只需要按照构建输入传给模型，执行forward就可以得到输出啦。一个简单的示例如下：
// Create a vector of inputs.
std::vector<torch::jit::IValue> inputs;
inputs.push_back(torch::ones({1, 3, 224, 224}));
// Execute the model and turn its output into a tensor.
at::Tensor output = module.forward(inputs).toTensor();
std::cout << output.slice(/*dim=*/1, /*start=*/0, /*end=*/5) << '\n';前两行创建一个torch::jit::IValue的向量，并添加单个输入. 使用torch::ones()创建输入张量，等效于C ++ API中的torch.ones。然后，运行script::Module的forward方法，通过调用toTensor()将返回的IValue值转换为张量。C++对torch的各种操作还是比较友好的，通过torch::或者后加_的方法都可以找到对应实现，例如
torch::tensor(input_list[j]).to(at::kLong).resize_({batch, 128}).clone()
//torch::tensor对应pytorch的torch.tensor; at::kLong对应torch.int64;resize_对应resize最后check一下确保c++端的输出和pytorch是一致的就大功告成啦~
踩了无数坑，薅掉了无数头发，很多东西也是自己一点点摸索的，如果有错误欢迎指正！
参考资料：
PyTorch C++ API - PyTorch master document
Torch Script - PyTorch master documentation
文章地址：
https://pytorch.org/cppdocs/
https://pytorch.org/tutorials/advanced/cpp_export.html
本文仅做学术分享，如有侵权，请联系删文。
后台回复：即可下载国外大学沉淀数年3D Vison精品课件
后台回复：即可下载3D视觉领域经典书籍pdf
后台回复：即可学习3D视觉领域精品课程
1.面向自动驾驶领域的多传感器数据融合技术
2.面向自动驾驶领域的3D点云目标检测全栈学习路线！(单模态+多模态/数据+代码)3.彻底搞透视觉三维重建：原理剖析、代码讲解、及优化改进4.国内首个面向工业级实战的点云处理课程5.激光-视觉-IMU-GPS融合SLAM算法梳理和代码讲解6.彻底搞懂视觉-惯性SLAM：基于VINS-Fusion正式开课啦7.彻底搞懂基于LOAM框架的3D激光SLAM: 源码剖析到算法优化8.彻底剖析室内、室外激光SLAM关键算法原理、代码和实战(cartographer+LOAM +LIO-SAM)
9.从零搭建一套结构光3D重建系统[理论+源码+实践]
10.单目深度估计方法：算法梳理与代码实现
11.自动驾驶中的深度学习模型部署实战
12.相机模型与标定(单目+双目+鱼眼）
13.重磅！四旋翼飞行器：算法与实战
14.ROS2从入门到精通：理论与实战
15.国内首个3D缺陷检测教程：理论、源码与实战
扫码添加小助手微信，可申请加入3D视觉工坊-学术论文写作与投稿 微信交流群，旨在
也可申请加入我们的细分方向交流群，目前主要有、、、、、等微信群，请扫描下面微信号加群，备注：”研究方向+学校/公司+昵称“，例如：”3D视觉 + 上海交大 + 静静“。请按照格式备注，否则不予通过。添加成功后会根据研究方向邀请进去相关微信群。也请联系。
▲长按加微信群或投稿
▲长按关注公众号
：针对3D视觉领域的五个方面进行深耕，更有各类大厂的算法工程人员进行技术指导。与此同时，星球将联合知名企业发布3D视觉相关算法开发岗位以及项目对接信息，打造成集技术与就业为一体的铁杆粉丝聚集区，近4000星球成员为创造更好的AI世界共同进步，知识星球入口：
学习3D视觉核心技术，扫描查看介绍，3天内无条件退款
圈里有高质量教程资料、答疑解惑、助你高效解决问题