大约一年前,PyTorch开发人员引入了TorchScript社区,该工具可让您通过几次可嵌入C ++系统的鼠标单击来从python的管道中创建可转让的解决方案。 下面,我分享使用它的经验,并尝试描述这条道路上遇到的陷阱。 我将特别注意Windows上该项目的实施,因为尽管ML的研究通常是在Ubuntu上完成的,但是最终的解决方案通常(突然!)在“ Windows”下是必需的。
可以在GitHub上的存储库中找到用于导出模型和使用该模型的C ++项目的示例代码。

PyTorch开发人员并未上当。 借助新工具,您可以在几个工作日内以更快的速度将PyTorch中的研究项目转换为嵌入C ++系统中的代码。
TorchScript出现在PyTorch 1.0版中,并且继续发展和变化。 如果一年前的第一个版本充满了漏洞并且更具实验性,那么至少在第二点上的当前版本1.3会明显不同:您不能再将其称为实验性的,它非常适合实际使用。 我会专注于她。
TorchScript的核心是其自己的类似Python的语言的独立(无Python)编译器,以及将以Python和PyTorch编写的程序转换为Python的工具,用于保存和加载生成的模块的方法以及用于在C ++中使用它们的库。 要工作,您必须向项目中添加几个DLL,总重约为70MB(对于Windows)才能在CPU上工作,而对于GPU版本则需要300MB。 TorchScript支持PyTorch的大多数功能以及python语言的主要功能。 但是第三方库,例如OpenCV或NumPy,将不得不被遗忘。 幸运的是,NumPy的许多功能在PyTorch中都有一个类似物。
在TorchScript上将管道转换为PyTorch模型
TorchScript提供了两种将Python代码转换为其内部格式的方法:跟踪和脚本编制(跟踪和脚本编制)。 为什么两个? 不,很明显,当然,两个比一个要好...

但是,就这些方法而言,与众所周知的格言一样,结果是左右偏差:两者都更糟。 好吧,世界并不完美。 仅在特定情况下,您需要选择更合适的一种。
跟踪方法非常简单。 采集数据样本(通常由随机数初始化),然后将其发送给我们感兴趣的类的函数或方法,然后PyTorch以与训练神经网络通常相同的方式构造和存储计算图。 瞧-脚本已经准备好:
import torch import torchvision model = torchvision.models.resnet34(pretrained = True) model.eval() sample = torch.rand(1, 3, 224, 224) scripted_model = torch.jit.trace(model, sample)
上面的示例生成ScriptModule类的对象。 可以保存
scripted_model.save('my_script.pth')
然后将其加载到C ++程序 (在下面进行更多介绍)或Python代码而不是原始对象中:
使用已保存的模型的示例Python代码 import cv2 from torchvision.transforms import Compose, ToTensor, Normalize transforms = Compose([ToTensor(), Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])]) img = cv2.resize(cv2.imread('pics/cat.jpg'), (224,224)) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) x = transforms(img).unsqueeze(0)
tensor(282) tensor(12.8130, grad_fn=<SelectBackward>)
生成的ScriptModule
可以出现在nn.Module
常用的任何位置。
以描述的方式,您可以跟踪nn.Module
类和函数的实例(在后一种情况下,将torch._C.Function
类的实例)。
这种方法(跟踪)具有一个重要的优点:这样,您几乎可以转换任何不使用外部库的Python代码。 但是有一个同样重要的缺点:对于任何分支,仅会记住在测试数据上执行的那个分支:
def my_abs(x): if x.max() >= 0: return x else: return -x my_abs_traced = torch.jit.trace(my_abs, torch.tensor(0)) print(my_abs_traced(torch.tensor(1)), my_abs_traced(torch.tensor(-1)))
c:\miniconda3\lib\site-packages\ipykernel_launcher.py:2: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs! tensor(1) tensor(-1)
糟糕! 这似乎不是我们想要的,对吗? 最好至少发出一条警告消息(TracerWarning)。 值得关注此类消息。
这是我们的第二种方法-脚本:
my_abs_script = torch.jit.script(my_abs) print(my_abs_script(torch.tensor(1)), my_abs_script(torch.tensor(-1)))
tensor(1) tensor(1)
万岁,收到了预期的结果! 脚本递归分析Python代码,并将其转换为自己语言的代码。 在输出中,我们还获得ScriptModule
类(对于模块)或torch._C.Function
(对于函数)。 似乎就是幸福! 但是出现了另一个问题:与Python不同,TorchScript的内部语言是强类型的。 每个变量的类型由第一个赋值确定,默认情况下,函数参数的类型为Tensor
。 因此,例如,一个熟悉的模式
def my_func(x): y = None if x.max() > 0: y = x return y my_func = torch.jit.script(my_func)
跟踪将失败。
跟踪错误看起来像这样 RuntimeError Traceback (most recent call last) <ipython-input-9-25414183a687> in <module>() ----> 1 my_func = torch.jit.script(my_func) d:\programming\3rd_party\pytorch\pytorch_ovod_1.3.0a0_de394b6\torch\jit\__init__.py in script(obj, optimize, _frames_up, _rcb) 1224 if _rcb is None: 1225 _rcb = _gen_rcb(obj, _frames_up) -> 1226 fn = torch._C._jit_script_compile(qualified_name, ast, _rcb, get_default_args(obj)) 1227 # Forward docstrings 1228 fn.__doc__ = obj.__doc__ RuntimeError: Variable 'y' previously has type None but is now being assigned to a value of type Tensor : at <ipython-input-8-75677614fca6>:4:8 def my_func(x): y = None if x.max() > 0: y = x ~ <--- HERE return y
值得注意的是,尽管在torch.jit.script
时发生错误,但在脚本代码中也指出了导致该错误的位置。
常数后的偶数点开始起作用:
def my_func(x): if x.max() > 0: y = 1.25 else: y = 0 return y my_func = torch.jit.script(my_func)
会给出一个错误 RuntimeError Traceback (most recent call last) <ipython-input-10-0a5f18586763> in <module>() 5 y = 0 6 return y ----> 7 my_func = torch.jit.script(my_func) d:\programming\3rd_party\pytorch\pytorch_ovod_1.3.0a0_de394b6\torch\jit\__init__.py in script(obj, optimize, _frames_up, _rcb) 1224 if _rcb is None: 1225 _rcb = _gen_rcb(obj, _frames_up) -> 1226 fn = torch._C._jit_script_compile(qualified_name, ast, _rcb, get_default_args(obj)) 1227 # Forward docstrings 1228 fn.__doc__ = obj.__doc__ d:\programming\3rd_party\pytorch\pytorch_ovod_1.3.0a0_de394b6\torch\jit\__init__.py in _rcb(name) 1240 # closure rcb fails 1241 result = closure_rcb(name) -> 1242 if result: 1243 return result 1244 return stack_rcb(name) RuntimeError: bool value of Tensor with more than one value is ambiguous
因为必须写的不是0
,而是0.
所以两个分支的类型是相同的! 您知道,使用python宠坏了!
这只是您需要对python代码进行更改的列表的开头,以便可以将其成功转换为TorchScript模块。 稍后,我将更详细地列出最典型的情况。 原则上,这里没有火箭科学,您的代码可以相应地更正。 但是大多数情况下,我不想修复第三方模块,包括torchvision
标准模块,并且通常它们通常不适合编写脚本。
幸运的是,两种技术可以结合使用:正在编写脚本的正在编写脚本,正在未编写脚本的正在跟踪:
class MyModule(torch.nn.Module): def __init__(self): super(MyModule, self).__init__() self.resnet = torchvision.models.resnet34(pretrained = True)
在上面的示例中,跟踪用于包含一个模块,该模块在没有足够跟踪且需要脚本的模块中无法编写脚本。 情况相反。 例如,如果我们需要将模型上传到ONNX,则使用跟踪。 但是跟踪模型可能包含TorchScript函数,因此可以在其中实现需要分支和循环的逻辑! 在torch.onnx的官方文档中给出了一个示例。
官方文档和torch.jit
中更详细地介绍了PyTorch提供的用于创建TorchScript模块的功能。 特别是,我没有提到以调试器形式使用torch.jit.trace
和torch.jit.script
的便捷方法,这涉及到调试脚本代码的特殊性。 这以及更多内容在文档中。
我们将模型包含在C ++项目中
不幸的是, 官方文档仅限于“添加使用torch.ones
生成的2张量”形式的示例。 我准备了一个更接近现实的项目示例, 该项目将图像从OpenCV发送到神经网络,并以响应张量,变量元组和带有分割结果的图像的形式接收结果。
为了使该示例生效,您需要使用ResNet34保存分类脚本,并使用DeepLabV3进行分段。 要准备这些脚本,您需要运行此jupyter notepad 。
我们需要torchlib
库。 您可以通过几种方式获得它:
- 如果您已经使用
pip install
PyTorch,则可以在Python目录中找到它: <Miniconda3>\Lib\site-packages\torch
- 如果您有从源代码编译的PyTorch,则它位于:
<My Pytorch repo>\build\lib.win-amd64-3.6\torch
; - 最后,您可以通过选择Language = C ++从pytorch.org单独下载该库 ,并解压缩存档。
C ++代码非常简单。 有必要:
- 包含头文件
#include <torch/script.h>
- 下载型号
torch::jit::script::Module module = torch::jit::load("../resnet34_infer.pth");
- 准备数据
torch::Tensor tensor = torch::from_blob(img.data, { img.rows, img.cols, 3 }, torch::kByte);
- 呼叫
forward
功能并获得结果
auto output = module.forward( { tensor } )
- 从结果中获取数据。 如何执行此操作取决于神经网络返回的内容。 顺便说一句,在一般情况下,它也不仅可以接受一张图片,因此最好看一下整个示例的源代码 ,有不同的选择。 例如,要从float类型的一维张量获取数据:
float* data = static_cast<float*>(output.toTensor().data_ptr());
- 还有另一种微妙之处。 不要忘记在代码中插入
with torch.no_grad()
的类似物with torch.no_grad()
在计算和存储不需要的梯度时浪费资源。 不幸的是,该命令不能包含在脚本中,因此您必须将其添加到C ++代码中:
torch::NoGradGuard no_grad;
官方指南中介绍了如何使用CMake构建项目。 但是那里没有公开Visual Studio中项目的主题,因此我将对其进行详细描述。 您将必须手动调整项目设置:
- 我在Visual Studio 2017上进行了测试。我不能说其他版本。
- 必须安装v14.11工具集v141(在VS安装程序中选中
"VC++ 2017 version 15.4 v14.11 toolset"
)。 - 该平台必须为
x64
。 - 在
General → Platform Toolset
v141(Visual Studio 2017)
General → Platform Toolset
选择v141(Visual Studio 2017)
- 在
C/C++ → General → Additional Include Directories
添加<libtorch dir>\include
- 在
Linker → General → Additional Library Directories
添加<libtorch dir>\lib
- 在
Linker → Input → Additional Dependencies
项中添加torch.lib; c10.lib
torch.lib; c10.lib
。 在Internet上,他们写道可能仍然需要caffe2.lib
,并且对于GPU和<libtorch dir>\lib
,但在当前版本中,添加这两个库对我来说已经足够。 也许这是过时的信息。 - 他们还写道,您需要设置
C/C++ → Language → Conformance Mode
= No
,但是我没有看到区别。
另外,不应在项目中声明__cplusplus
变量。 尝试添加 /Zc:__cplusplus
将在ivalue.h
文件中导致编译错误。
在附件项目中,路径设置(不仅包括TorchLib的路径设置,还包括OpenCV和CUDA的路径设置)都在props文件中取出,在进行组装之前,您需要根据本地配置在此处注册它们。 实际上,仅此而已。
还有什么要记住的
如果所描述的过程对您来说似乎太简单了,那么您的直觉并不会欺骗您。 为了将用Python编写的PyTorch模型转换为TorchScript,需要考虑许多细微差别。 我将在下面列出我必须面对的内容。 我已经提到过一些,但是我重复一遍将所有内容收集在一个地方。

- 默认情况下,传递给函数的变量类型为Tensor。 如果在某些情况下(非常频繁),这是不可接受的,则必须使用MyPy样式类型注释手动声明类型,如下所示:
def calc_letter_statistics(self, cls_preds: List[Tensor], cls_thresh: float)->Tuple[int, Tuple[Tensor, Tensor, Tensor]]
左右:
def calc_letter_statistics(self, cls_preds, cls_thresh):
- 变量是强类型的,并且如果未明确指定,则类型由第一个赋值确定。 形式
x=[]; for ...: x.append(y)
熟悉的构造x=[]; for ...: x.append(y)
x=[]; for ...: x.append(y)
必须进行编辑,因为 在分配[]
编译器无法确定列表中的类型。 因此,您必须明确指定类型,例如:
from typing import List x: List[float] = []
或(另一个“例如”)
from torch import Tensor from typing import Dict, Tuple, List x: Dict[int: Tuple[float, List[Tensor], List[List[int]]]] = {}
- 在上面的示例中,由于这些名称被缝在TorchScript代码中,因此需要导入这些名称。 替代的看似合法的方法
import torch import typing x: typing.List[torch.Tensor] = []
会导致输入类型未知的构造函数 。编写脚本时出现列表错误
x = None if smth: x = torch.tensor([1,2,3])
有两种选择。 或两次分配张量(尺寸不同的事实并不可怕):
x = torch.tensor(0) if smth: x = torch.tensor([1,2,3])
并且不要忘记寻找更换后会破裂的东西。 或者尝试诚实地写:
x: Optional[Tensor] = None if smth: x = torch.tensor([1,2,3])
但是如果在预期的张量处进一步使用x
,我们很可能会得到一个错误: 期望参数'x'的类型为'Tensor',但找到的类型为'Optional [Tensor]'。
别忘了在第一次分配期间写x=0.
如果变量x
必须为float
类型,则代替通常的x=0
等。
如果在某个地方我们通过x = torch.Tensor(...)
使用了张量的老式初始化,则您必须将其x = torch.Tensor(...)
,并用带有小写字母x = torch.tensor(...)
的较新版本替换它。 否则,它将在脚本编写过程中飞行: 未知的内置操作:aten :: Tensor。 以下是一些建议:aten ::张量 。 看来他们甚至可以解释问题所在,并且很清楚需要做什么。 但是,很显然您是否已经知道正确的答案。
该代码在调用torch.jit.script
的模块的上下文中编写脚本。 因此,如果在脚本化类或函数的肠道中某处(例如,使用math.pow
,则必须将import math
添加到编译模块中。 而且最好在声明该类的脚本上编写脚本:使用@torch.jit.script
装饰器,或在其旁边声明一个使ScriptModule脱离其功能的附加函数。 否则,当我们尝试从其中显然进行了math
导入的模块中编译类时,会收到未定义值的数学错误消息。
如果在某处您具有my_tensor[my_tensor < 10] = 0
或类似形式的构造,则在编写脚本时会收到一个神秘的错误:
*aten::index_put_(Tensor(a!) self, Tensor?[] indices, Tensor values, bool accumulate=False) -> (Tensor(a!)):* *Expected a value of type 'Tensor' for argument 'values' but instead found type 'int'.* *aten::index_put_(Tensor(a!) self, Tensor[] indices, Tensor values, bool accumulate=False) -> (Tensor(a!)):* *Expected a value of type 'List[Tensor]' for argument 'indices' but instead found type 'List[Optional[Tensor]]'.*
您需要用张量替换数字: my_tensor[my_tensor < 10] = torch.tensor(0.).to(my_tensor.device)
。 并且不要忘记a)关于my_tensor
类型与创建的张量(在这种情况下为float)的对应关系,以及b)关于.to(my_tensor.device)
。 如果您忘记了第二个,所有内容都将被编写脚本,但是已经在使用GPU的过程中,您将感到不安,这看起来像是CUDA错误的字眼:遇到了非法的内存访问 ,而没有指出发生错误的地方!
不要忘记默认情况下nn.Module
以及相应的torchvision模型是在“训练模式”下创建的(您不会相信,但事实证明确实存在这种模式 )。 在这种情况下,将使用Dropout和火车模式下的其他技巧,这些技巧可能会中断跟踪或在执行时导致结果不足。 切记在编写脚本或跟踪之前调用model.eval()
。
对于函数和普通类,您需要为nn.Module-一个实例编写类型脚本
尝试使用脚本方法访问全局变量
cls_thresh = 0.3 class MyModule(torch.nn.Module): ... x = r < cls_thresh ...
将导致形式为'float'的python值形式的脚本错误, 不能用作value 。 必须在构造函数中使变量成为属性:
cls_thresh = 0.3 class MyModule(torch.nn.Module): def __init__(self): ... self.cls_thresh = cls_thresh ... x = r < self.cls_thresh ...
- 如果将class属性用作slice参数,则会产生另一个微妙的问题:
class FPN(nn.Module): def __init__(self, block, num_blocks, num_layers =5): ... self.num_layers = num_layers def forward(self, x): ... return (p3, p4, p5, p6, p7)[:self.num_layers]
导致脚本错误的元组切片索引必须是整数常量 。 有必要指出num_layers属性是恒定的,并且不会改变:
class FPN(nn.Module): num_layers: torch.jit.Final[int] def __init__(self, block, num_blocks, num_layers =5): ...
- 在某些情况下,张量通常可以正常拟合,您需要显式传递数字:
xx1 = x1.clamp(min=x1[i])
在脚本中Expected a value of type 'Optional[number]' for argument 'min' but instead found type 'Tensor'.
时,抛出错误Expected a value of type 'Optional[number]' for argument 'min' but instead found type 'Tensor'.
。 好吧,从错误消息中可以很清楚地知道该怎么做:
xx1 = x1.clamp(min=x1[i].item())
跟踪时出现上述问题。 由于它们的存在,通常无法在TorchScript中简单地编译现成的解决方案,并且您必须长时间按摩源代码(如果源代码适合于编辑),或者使用跟踪。 但是痕迹有其细微差别:
tensor_a.to(tensor_b.device)
张量加载到的设备在跟踪时是固定的,在执行期间不会更改。 通过将张量声明为nn.Module
类型的nn.Module
成员,可以部分克服此问题。 然后,在加载模型时,它将启动到torch.jit.load
函数中指定的设备上。
结语
当然,以上所有都会产生问题。 但是,TorchScript允许您将模型本身以及提供预处理和后处理的Python代码作为一个整体组合并发送到解决方案。 是的,尽管存在上述困难,但准备解决方案进行编译所需的时间却比创建解决方案的成本少得多,但是PyTorch在这里提供了巨大的优势,因此值得一试。
