TorchScript 简介#
参考：hybridize & Intro_to_TorchScript_tutorial
TorchScript 允许用户使用纯命令式编程进行开发和调试，同时能够将大多数程序转换为符号式程序，以便在需要产品级计算性能和部署时使用。
TorchScript 是 PyTorch 模型的中间表示形式（Module 的子类），可以在 C++ 等高性能环境中运行。
PyTorch 模型创建的基础知识#
从定义简单的 Module 开始。Module 是 PyTorch 中的基本组成单元。它包含：
构造函数，它为调用模块做准备
一组
Parameters和子Module。它们由构造函数初始化，module可以在调用期间使用forward函数。这是调用模块时运行的代码。
简单示例：
import torch # This is all you need to use both PyTorch and TorchScript!
print(torch.__version__)
torch.manual_seed(191009) # set the seed for reproducibility
2.5.1
<torch._C.Generator at 0x7f1c2e15a130>
PyTorch 模型编写基础#
从简单的 torch.nn.Module 定义开始。torch.nn.Module 是 PyTorch 中的基本组成单元。它包含：
构造函数，用于准备模块以供调用
一组
Parameters和子Modules。这些由构造函数初始化，并在模块调用期间使用forward函数。这是在模块被调用时执行的代码
小例子：
class MyCell(torch.nn.Module):
def __init__(self):
super().__init__()
def forward(self, x, h):
new_h = torch.tanh(x + h)
return new_h, new_h
my_cell = MyCell()
x = torch.rand(3, 4)
h = torch.rand(3, 4)
print(my_cell(x, h))
(tensor([[0.8951, 0.9477, 0.8354, 0.7361],
[0.8297, 0.7231, 0.8003, 0.8989],
[0.7820, 0.6564, 0.7702, 0.2425]]), tensor([[0.8951, 0.9477, 0.8354, 0.7361],
[0.8297, 0.7231, 0.8003, 0.8989],
[0.7820, 0.6564, 0.7702, 0.2425]]))
上述过程如下：
创建了继承自
torch.nn.Module的类定义了构造函数。构造函数没有做太多事情，只是调用了
super的构造函数定义了
forward函数，该函数接受两个输入并返回两个输出。forward函数的实际内容并不重要，但它类似于假的 RNN 单元——也就是说，它是在循环中应用的函数
实例化模块，并创建了 x 和 h，它们只是随机值的 3x4 矩阵。然后使用 my_cell(x, h) 调用了该单元。这反过来调用了我们的 forward 函数。
做一些更有趣的事情：
class MyCell(torch.nn.Module):
def __init__(self):
super(MyCell, self).__init__()
self.linear = torch.nn.Linear(4, 4)
def forward(self, x, h):
new_h = torch.tanh(self.linear(x) + h)
return new_h, new_h
my_cell = MyCell()
print(my_cell)
print(my_cell(x, h))
MyCell(
(linear): Linear(in_features=4, out_features=4, bias=True)
)
(tensor([[-0.0523, 0.7750, 0.1299, 0.7710],
[-0.0793, 0.3363, -0.1186, 0.6945],
[-0.0596, 0.4099, 0.3653, 0.3728]], grad_fn=<TanhBackward0>), tensor([[-0.0523, 0.7750, 0.1299, 0.7710],
[-0.0793, 0.3363, -0.1186, 0.6945],
[-0.0596, 0.4099, 0.3653, 0.3728]], grad_fn=<TanhBackward0>))
重新定义了模块 MyCell，但这次我们添加了 self.linear 属性，并在 forward 函数中调用了 self.linear。
这里到底发生了什么？torch.nn.Linear 是 PyTorch 标准库中的 Module。就像 MyCell 一样，它可以通过调用语法来调用。正在构建 Module 的层次结构。
在 Module 上使用 print 将给出 Module 子类层次结构的视觉表示。
通过以这种方式组合 Module，可以简洁且可读地使用可重用的组件来编写模型。
可能已经注意到输出中的 grad_fn。这是 PyTorch 自动微分方法的细节，称为 autograd。简而言之，该系统允许通过可能复杂的程序计算导数。这种设计允许在模型编写中具有极大的灵活性。
现在让我们来探讨这种灵活性：
class MyDecisionGate(torch.nn.Module):
def forward(self, x):
if x.sum() > 0:
return x
else:
return -x
class MyCell(torch.nn.Module):
def __init__(self):
super(MyCell, self).__init__()
self.dg = MyDecisionGate()
self.linear = torch.nn.Linear(4, 4)
def forward(self, x, h):
new_h = torch.tanh(self.dg(self.linear(x)) + h)
return new_h, new_h
my_cell = MyCell()
print(my_cell)
print(my_cell(x, h))
MyCell(
(dg): MyDecisionGate()
(linear): Linear(in_features=4, out_features=4, bias=True)
)
(tensor([[0.5886, 0.6279, 0.6672, 0.8669],
[0.6514, 0.0418, 0.6567, 0.6881],
[0.4862, 0.2474, 0.6690, 0.5180]], grad_fn=<TanhBackward0>), tensor([[0.5886, 0.6279, 0.6672, 0.8669],
[0.6514, 0.0418, 0.6567, 0.6881],
[0.4862, 0.2474, 0.6690, 0.5180]], grad_fn=<TanhBackward0>))
再次重新定义了 MyCell 类，但这次定义了 MyDecisionGate。这个模块使用了 控制流。控制流包括循环和 if 语句等。
许多框架采用的方法是给定完整程序表示来计算符号导数。然而，在 PyTorch 中，使用梯度带（gradient tape）。在操作发生时记录它们，并在反向传播时重放它们以计算导数。通过这种方式，框架不需要显式地为语言中的所有构造定义导数。
TorchScript 基础#
现在让我们以正在运行的示例为例，看看如何应用 TorchScript。
简而言之，TorchScript 提供了工具来捕获模型的定义，即使在 PyTorch 的灵活和动态特性下也是如此。让我们从检查我们称之为 追踪 的内容开始。
追踪 Module#
class MyCell(torch.nn.Module):
def __init__(self):
super(MyCell, self).__init__()
self.linear = torch.nn.Linear(4, 4)
def forward(self, x, h):
new_h = torch.tanh(self.linear(x) + h)
return new_h, new_h
my_cell = MyCell()
x, h = torch.rand(3, 4), torch.rand(3, 4)
traced_cell = torch.jit.trace(my_cell, (x, h))
print(traced_cell)
traced_cell(x, h)
MyCell(
original_name=MyCell
(linear): Linear(original_name=Linear)
)
(tensor([[ 0.4021, 0.7634, 0.6591, -0.3496],
[ 0.7640, 0.6614, 0.2497, -0.1470],
[ 0.4033, 0.2158, 0.3466, -0.4910]], grad_fn=<TanhBackward0>),
tensor([[ 0.4021, 0.7634, 0.6591, -0.3496],
[ 0.7640, 0.6614, 0.2497, -0.1470],
[ 0.4033, 0.2158, 0.3466, -0.4910]], grad_fn=<TanhBackward0>))
回溯了一点，采用了 MyCell 类的第二个版本。和之前一样，实例化了它，但这次调用了 torch.jit.trace，传入了 Module，并传入了网络可能看到的 示例输入。
这到底做了什么？它调用了 Module，记录了运行 Module 时发生的操作，并创建了 torch.jit.ScriptModule 的实例（其中 TracedModule 是一个实例）。
TorchScript 在其中间表示（或 IR）中记录其定义，这在深度学习中通常被称为 计算图。可以使用 .graph 属性检查计算图：
print(traced_cell.graph)
graph(%self.1 : __torch__.MyCell,
%x : Float(3, 4, strides=[4, 1], requires_grad=0, device=cpu),
%h : Float(3, 4, strides=[4, 1], requires_grad=0, device=cpu)):
%linear : __torch__.torch.nn.modules.linear.Linear = prim::GetAttr[name="linear"](%self.1)
%20 : Tensor = prim::CallMethod[name="forward"](%linear, %x)
%11 : int = prim::Constant[value=1]() # /tmp/ipykernel_2202827/260609686.py:7:0
%12 : Float(3, 4, strides=[4, 1], requires_grad=1, device=cpu) = aten::add(%20, %h, %11) # /tmp/ipykernel_2202827/260609686.py:7:0
%13 : Float(3, 4, strides=[4, 1], requires_grad=1, device=cpu) = aten::tanh(%12) # /tmp/ipykernel_2202827/260609686.py:7:0
%14 : (Float(3, 4, strides=[4, 1], requires_grad=1, device=cpu), Float(3, 4, strides=[4, 1], requires_grad=1, device=cpu)) = prim::TupleConstruct(%13, %13)
return (%14)
然而，这是非常低级的表示，计算图中包含的大部分信息对最终用户来说是没有用的。相反，可以使用 .code 属性来给出代码的 Python 语法解释：
print(traced_cell.code)
def forward(self,
x: Tensor,
h: Tensor) -> Tuple[Tensor, Tensor]:
linear = self.linear
_0 = torch.tanh(torch.add((linear).forward(x, ), h))
return (_0, _0)
那么 为什么 我们要做所有这些呢？有几个原因：
TorchScript 代码可以在其自己的解释器中调用，该解释器基本上是受限的 Python 解释器。这个解释器不会获取全局解释器锁（GIL），因此可以在同一实例上同时处理许多请求。
这种格式允许整个模型保存到磁盘并在另一个环境中加载它，例如在用 Python 以外的语言编写的服务器中。
TorchScript 提供了一种表示形式，可以在其中对代码进行编译器优化以提供更高效的执行。
TorchScript 允许与许多需要比单个算子更广泛的程序视图的后端/设备运行时进行接口。
我们可以看到，调用 traced_cell 产生的结果与 Python 模块相同：
print(my_cell(x, h))
print(traced_cell(x, h))
(tensor([[ 0.4021, 0.7634, 0.6591, -0.3496],
[ 0.7640, 0.6614, 0.2497, -0.1470],
[ 0.4033, 0.2158, 0.3466, -0.4910]], grad_fn=<TanhBackward0>), tensor([[ 0.4021, 0.7634, 0.6591, -0.3496],
[ 0.7640, 0.6614, 0.2497, -0.1470],
[ 0.4033, 0.2158, 0.3466, -0.4910]], grad_fn=<TanhBackward0>))
(tensor([[ 0.4021, 0.7634, 0.6591, -0.3496],
[ 0.7640, 0.6614, 0.2497, -0.1470],
[ 0.4033, 0.2158, 0.3466, -0.4910]], grad_fn=<TanhBackward0>), tensor([[ 0.4021, 0.7634, 0.6591, -0.3496],
[ 0.7640, 0.6614, 0.2497, -0.1470],
[ 0.4033, 0.2158, 0.3466, -0.4910]], grad_fn=<TanhBackward0>))
使用脚本转换模块#
使用模块的第二个版本，而不是带有控制流子模块的版本，是有原因的。现在让我们来检查一下：
class MyDecisionGate(torch.nn.Module):
def forward(self, x):
if x.sum() > 0:
return x
else:
return -x
class MyCell(torch.nn.Module):
def __init__(self, dg):
super(MyCell, self).__init__()
self.dg = dg
self.linear = torch.nn.Linear(4, 4)
def forward(self, x, h):
new_h = torch.tanh(self.dg(self.linear(x)) + h)
return new_h, new_h
my_cell = MyCell(MyDecisionGate())
traced_cell = torch.jit.trace(my_cell, (x, h))
print(traced_cell.dg.code)
print(traced_cell.code)
def forward(self,
argument_1: Tensor) -> Tensor:
return torch.neg(argument_1)
def forward(self,
x: Tensor,
h: Tensor) -> Tuple[Tensor, Tensor]:
dg = self.dg
linear = self.linear
_0 = torch.add((dg).forward((linear).forward(x, ), ), h)
_1 = torch.tanh(_0)
return (_1, _1)
/tmp/ipykernel_2202827/4234398751.py:3: 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!
if x.sum() > 0:
查看 .code 输出，我们可以看到 if-else 分支无处可寻！为什么？追踪确实做了我们所说的：运行代码，记录 发生的 运算，并构建 ScriptModule 来执行这些操作。不幸的是，像控制流这样的东西被删除了。
如何才能在 TorchScript 中忠实地表示这个模块？脚本编译器 可以直接分析您的 Python 源代码以将其转换为 TorchScript。使用脚本编译器来转换 MyDecisionGate：
scripted_gate = torch.jit.script(MyDecisionGate())
my_cell = MyCell(scripted_gate)
scripted_cell = torch.jit.script(my_cell)
print(scripted_gate.code)
print(scripted_cell.code)
def forward(self,
x: Tensor) -> Tensor:
if bool(torch.gt(torch.sum(x), 0)):
_0 = x
else:
_0 = torch.neg(x)
return _0
def forward(self,
x: Tensor,
h: Tensor) -> Tuple[Tensor, Tensor]:
dg = self.dg
linear = self.linear
_0 = torch.add((dg).forward((linear).forward(x, ), ), h)
new_h = torch.tanh(_0)
return (new_h, new_h)
万岁！我们现在已经在 TorchScript 中忠实地捕获了程序的行为。现在让我们尝试运行程序：
# New inputs
x, h = torch.rand(3, 4), torch.rand(3, 4)
print(scripted_cell(x, h))
(tensor([[0.1675, 0.9159, 0.5259, 0.8139],
[0.8376, 0.9887, 0.7665, 0.7780],
[0.4424, 0.9728, 0.6343, 0.8376]], grad_fn=<TanhBackward0>), tensor([[0.1675, 0.9159, 0.5259, 0.8139],
[0.8376, 0.9887, 0.7665, 0.7780],
[0.4424, 0.9728, 0.6343, 0.8376]], grad_fn=<TanhBackward0>))
混合脚本和追踪#
有些情况需要使用追踪而不是脚本（例如，模块有许多基于常量 Python 值的架构决策，我们不希望这些决策出现在 TorchScript 中）。在这种情况下，脚本可以与追踪组合使用：torch.jit.script 将内联追踪模块的代码，而追踪将内联脚本模块的代码。
第一个情况的示例：
class MyRNNLoop(torch.nn.Module):
def __init__(self):
super(MyRNNLoop, self).__init__()
self.cell = torch.jit.trace(MyCell(scripted_gate), (x, h))
def forward(self, xs):
h, y = torch.zeros(3, 4), torch.zeros(3, 4)
for i in range(xs.size(0)):
y, h = self.cell(xs[i], h)
return y, h
rnn_loop = torch.jit.script(MyRNNLoop())
print(rnn_loop.code)
def forward(self,
xs: Tensor) -> Tuple[Tensor, Tensor]:
h = torch.zeros([3, 4])
y = torch.zeros([3, 4])
y0 = y
h0 = h
for i in range(torch.size(xs, 0)):
cell = self.cell
_0 = (cell).forward(torch.select(xs, 0, i), h0, )
y1, h1, = _0
y0, h0 = y1, h1
return (y0, h0)
第二个情况的示例：
class WrapRNN(torch.nn.Module):
def __init__(self):
super(WrapRNN, self).__init__()
self.loop = torch.jit.script(MyRNNLoop())
def forward(self, xs):
y, h = self.loop(xs)
return torch.relu(y)
traced = torch.jit.trace(WrapRNN(), (torch.rand(10, 3, 4)))
print(traced.code)
def forward(self,
xs: Tensor) -> Tensor:
loop = self.loop
_0, y, = (loop).forward(xs, )
return torch.relu(y)
这样，脚本和追踪可以在需要时分别使用，并且可以一起使用。
保存和加载模型#
我们提供了 API 来将 TorchScript 模块保存到磁盘并从磁盘加载，使用存档格式。这种格式包括代码、参数、属性和调试信息，这意味着存档是模型的独立表示，可以在完全独立的进程中加载。让我们保存并加载我们包装的 RNN 模块：
traced.save('wrapped_rnn.pt')
loaded = torch.jit.load('wrapped_rnn.pt')
print(loaded)
print(loaded.code)
RecursiveScriptModule(
original_name=WrapRNN
(loop): RecursiveScriptModule(
original_name=MyRNNLoop
(cell): RecursiveScriptModule(
original_name=MyCell
(dg): RecursiveScriptModule(original_name=MyDecisionGate)
(linear): RecursiveScriptModule(original_name=Linear)
)
)
)
def forward(self,
xs: Tensor) -> Tensor:
loop = self.loop
_0, y, = (loop).forward(xs, )
return torch.relu(y)
正如您所见，序列化保留了模块层次结构和我们一直在检查的代码。该模型也可以加载，例如，加载到 C++ 中以进行无 Python 的执行。
进一步阅读#
我们已经完成了本教程！对于更深入的演示，请查看使用 TorchScript 转换机器翻译模型的 NeurIPS 演示：https://colab.research.google.com/drive/1HiICg6jRkBnr5hvK2-VnMi88Vi9pUzEJ