TorchScript 简介

参考：hybridize

TorchScript 允许用户使用纯命令式编程进行开发和调试，同时能够将大多数程序转换为符号式程序，以便在需要产品级计算性能和部署时使用。

TorchScript 是 PyTorch 模型的中间表示形式（Module 的子类），可以在 C++ 等高性能环境中运行。

PyTorch 模型创建的基础知识

从定义简单的 Module 开始。Module 是 PyTorch 中的基本组成单元。它包含：

• 构造函数，它为调用模块做准备

• 一组 Parameters 和子 Module。它们由构造函数初始化，module 可以在调用期间使用

• forward 函数。这是调用模块时运行的代码。

简单示例：

import torch


class MyCell(torch.nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x, h):
        new_h = torch.tanh(x + h)
        return new_h

my_cell = MyCell()
x = torch.rand(3, 4)
h = torch.rand(3, 4)
print(my_cell(x, h))

tensor([[0.7110, 0.9077, 0.6591, 0.5533],
        [0.7155, 0.7255, 0.7911, 0.6922],
        [0.4663, 0.5373, 0.4477, 0.8374]])

class MyCell(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = torch.nn.Linear(4, 4)

    def forward(self, x, h):
        new_h = torch.tanh(self.linear(x) + h)
        return 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.8500, 0.9233, 0.5103, 0.2327],
        [0.7223, 0.8732, 0.2950, 0.7843],
        [0.7316, 0.7307, 0.0477, 0.8260]], 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):
        super().__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

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.7368,  0.9277,  0.4715,  0.0517],
        [ 0.7383,  0.8358,  0.4346,  0.6980],
        [ 0.6292,  0.8111, -0.1753,  0.7719]], grad_fn=<TanhBackward0>)

TorchScript 基础

看看如何应用 TorchScript。简而言之，即使考虑到 PyTorch 的灵活性和动态特性，TorchScript 也提供了捕获模型定义的工具。让我们从研究所谓的 trace() 开始。

class MyCell(torch.nn.Module):
    def __init__(self):
        super().__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.2799,  0.6735,  0.3911,  0.2816],
         [-0.2376,  0.8031,  0.0191,  0.7671],
         [ 0.1377,  0.5180, -0.6820, -0.0965]], grad_fn=<TanhBackward0>),
 tensor([[ 0.2799,  0.6735,  0.3911,  0.2816],
         [-0.2376,  0.8031,  0.0191,  0.7671],
         [ 0.1377,  0.5180, -0.6820, -0.0965]], grad_fn=<TanhBackward0>))

跟踪函数并返回可执行文件或 ScriptFunction，它将使用即时编译进行优化。trace() 对于仅对 Tensor 以及 Tensor 的列表、字典和元组进行运算的代码是理想的。

type(traced_cell)

torch.jit._trace.TopLevelTracedModule

使用 torch.jit.trace() 和 torch.jit.trace_module()，你可以将现有的模块或 Python 函数转换为 ScriptFunction 或 ScriptModule。你必须提供样本输入，然后运行这个函数，记录所有张量上的运算。

• 独立函数的结果记录产生 ScriptFunction。

• nn.Module.forward() 的结果记录生成 ScriptModule。

该模块还包含原始模块所具有的任何参数。

警告

追踪只正确记录不依赖于数据的函数和模块（例如，对张量中的数据没有条件），并且没有任何未跟踪的外部依赖（例如，执行输入/输出或访问全局变量）。追踪只记录给定函数在给定张量上运行时所做的运算。因此，返回的 ScriptModule 将始终在任何输入上运行相同的跟踪图。当您的模块需要根据输入和/或模块状态运行不同的运算集时，这有一些重要的含义。例如

• 跟踪不会记录任何控制流，如 if 语句或循环。当这个控制流在您的模块中是恒定的时，这很好，并且它通常内联控制流决策。但有时控制流实际上是模型本身的一部分。例如，循环网络是一个循环（可能是动态的）输入序列的长度。

• 在返回的 ScriptModule 中，在 training 和 eval 模式中具有不同行为的运算将始终表现得好像它处于跟踪期间所处的模式，无论 ScriptModule 处于哪种模式。

在这样的情况下，跟踪是不合适的，而 torch.jit.script() 是更好的选择。如果您跟踪这样的模型，您可能会在后续模型的调用运算中得到不正确的结果。当执行可能导致产生错误跟踪的运算时，跟踪程序将尝试发出警告。

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_343239/2138254843.py:7:0
  %12 : Float(3, 4, strides=[4, 1], requires_grad=1, device=cpu) = aten::add(%20, %h, %11) # /tmp/ipykernel_343239/2138254843.py:7:0
  %13 : Float(3, 4, strides=[4, 1], requires_grad=1, device=cpu) = aten::tanh(%12) # /tmp/ipykernel_343239/2138254843.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 解释器。该解释器不获得全局解释器锁，因此可以在同一个实例上同时处理许多请求。

• 这种格式允许我们将整个模型保存到磁盘上，并将其加载到另一个环境中，例如在用 Python 以外的语言编写的服务器中。

• TorchScript 为我们提供了一种表示法，可以在其中对代码进行编译器优化，以提供更高效的执行。

• TorchScript 允许与许多后端/设备运行时交互，这些运行时需要比单个算子更广泛的程序视图。

可以看到，调用 traced_cell 会产生与 Python 模块相同的结果：

print(my_cell(x, h))
print(traced_cell(x, h))

(tensor([[ 0.2799,  0.6735,  0.3911,  0.2816],
        [-0.2376,  0.8031,  0.0191,  0.7671],
        [ 0.1377,  0.5180, -0.6820, -0.0965]], grad_fn=<TanhBackward0>), tensor([[ 0.2799,  0.6735,  0.3911,  0.2816],
        [-0.2376,  0.8031,  0.0191,  0.7671],
        [ 0.1377,  0.5180, -0.6820, -0.0965]], grad_fn=<TanhBackward0>))
(tensor([[ 0.2799,  0.6735,  0.3911,  0.2816],
        [-0.2376,  0.8031,  0.0191,  0.7671],
        [ 0.1377,  0.5180, -0.6820, -0.0965]],
       grad_fn=<DifferentiableGraphBackward>), tensor([[ 0.2799,  0.6735,  0.3911,  0.2816],
        [-0.2376,  0.8031,  0.0191,  0.7671],
        [ 0.1377,  0.5180, -0.6820, -0.0965]],
       grad_fn=<DifferentiableGraphBackward>))

使用脚本转换模块

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().__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_343239/3110858787.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 中忠实地表示这个模块？Torch 提供了脚本编译器，它直接分析您的 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)

x, h = torch.rand(3, 4), torch.rand(3, 4)
traced_cell(x, h)

(tensor([[-0.2027,  0.9097,  0.1824,  0.5036],
         [ 0.2261,  0.8793,  0.3904,  0.4577],
         [ 0.1615,  0.9411,  0.1274,  0.4738]], grad_fn=<TanhBackward0>),
 tensor([[-0.2027,  0.9097,  0.1824,  0.5036],
         [ 0.2261,  0.8793,  0.3904,  0.4577],
         [ 0.1615,  0.9411,  0.1274,  0.4738]], grad_fn=<TanhBackward0>))

混合脚本和跟踪

有些情况需要使用 trace() 而不是 script()（例如，模块有许多基于常量 Python 值的架构决策，而我们不希望这些常量 Python 值出现在 TorchScript 中）。在这种情况下，script() 可以与 trace() 组合在一起：torch.jit.script 将内联被跟踪模块的代码，而跟踪将内联被脚本模块的代码。

class MyRNNLoop(torch.nn.Module):
    def __init__(self):
        super().__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().__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)

这样，当情况需要脚本和跟踪时，就可以使用它们并将它们一起使用。

保存和加载模型

Torch 提供 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 执行。