【官方文档解读】torch.jit.script 的使用，并附上官方文档中的示例代码

由 OpenMMLab 的部署教程所述，对于模型中存在有控制条件的（如 if，for 等），需要用 torch.jit.script 而非采样默认的 torch.jit.trace 方法。本文则详细介绍了下官方文档中对 torch.jit.script 的解释和示例代码。

torch.jit.script

torch.jit.script 用于将函数或 nn.Module 编译为 TorchScript。

函数签名

torch.jit.script(obj, optimize=None, _frames_up=0, _rcb=None, example_inputs=None)

功能概述

将函数或 nn.Module 脚本化，会检查源代码，并使用 TorchScript 编译器将其编译为 TorchScript 代码，并返回一个 ScriptModule 或 ScriptFunction。TorchScript 是 Python 语言的一个子集，因此并不是所有的 Python 功能都能在其中使用，但我们提供了足够的功能来对张量进行计算和执行控制相关操作。完整指南请参阅 TorchScript 语言参考。

脚本化字典或列表会将其中的数据复制到一个 TorchScript 实例中，该实例可以在 Python 和 TorchScript 之间以零复制开销传递引用。

torch.jit.script 可以作为函数用于模块、函数、字典和列表，并可以作为装饰器 @torch.jit.script 用于 TorchScript 类和函数。

参数

obj（Callable、类或 nn.Module） – 要编译的 nn.Module、函数、类类型、字典或列表。
example_inputs（Union[List[Tuple], Dict[Callable, List[Tuple]], None]） – 提供示例输入以注释函数或 nn.Module 的参数。

返回值

如果 obj 是 nn.Module，脚本会返回一个 ScriptModule 对象。返回的 ScriptModule 将具有与原始 nn.Module 相同的子模块和参数集。如果 obj 是独立函数，将返回 ScriptFunction。如果 obj 是字典，则脚本返回 torch._C.ScriptDict 实例。如果 obj 是列表，则脚本返回 torch._C.ScriptList 实例。

脚本化函数

@torch.jit.script 装饰器通过编译函数体来构建 ScriptFunction。

示例（脚本化函数）：

import torch@torch.jit.script
def foo(x, y):if x.max() > y.max():r = xelse:r = yreturn rprint(type(foo))  # torch.jit.ScriptFunction# 以 Python 代码查看编译后的图
print(foo.code)# 使用 TorchScript 解释器调用函数
foo(torch.ones(2, 2), torch.ones(2, 2))

使用示例输入脚本化函数

示例输入可用于注释函数参数。

示例（脚本化前注释函数）：

import torchdef test_sum(a, b):return a + b# 注释参数为 int
scripted_fn = torch.jit.script(test_sum, example_inputs=[(3, 4)])print(type(scripted_fn))  # torch.jit.ScriptFunction# 以 Python 代码查看编译后的图
print(scripted_fn.code)# 使用 TorchScript 解释器调用函数
scripted_fn(20, 100)

脚本化 nn.Module

默认情况下，脚本化 nn.Module 会编译 forward 方法，并递归编译 forward 调用的任何方法、子模块和函数。如果 nn.Module 仅使用 TorchScript 支持的功能，则无需对原始模块代码进行任何更改。脚本将构建一个 ScriptModule，其中包含原始模块的属性、副本和方法。

示例（脚本化包含参数的简单模块）：

import torchclass MyModule(torch.nn.Module):def __init__(self, N, M):super().__init__()# 此参数将被复制到新的 ScriptModuleself.weight = torch.nn.Parameter(torch.rand(N, M))# 当使用此子模块时，它将被编译self.linear = torch.nn.Linear(N, M)def forward(self, input):output = self.weight.mv(input)# 这会调用 `nn.Linear` 模块的 `forward` 方法，从而在此处将 `self.linear` 子模块编译为 `ScriptModule`output = self.linear(output)return outputscripted_module = torch.jit.script(MyModule(2, 3))

示例（脚本化包含 traced 子模块的模块）：

import torch
import torch.nn as nn
import torch.nn.functional as Fclass MyModule(nn.Module):def __init__(self):super().__init__()# torch.jit.trace 生成一个 ScriptModule 的 conv1 和 conv2self.conv1 = torch.jit.trace(nn.Conv2d(1, 20, 5), torch.rand(1, 1, 16, 16))self.conv2 = torch.jit.trace(nn.Conv2d(20, 20, 5), torch.rand(1, 20, 16, 16))def forward(self, input):input = F.relu(self.conv1(input))input = F.relu(self.conv2(input))return inputscripted_module = torch.jit.script(MyModule())

要编译 forward 以外的方法（并递归编译它调用的任何内容），请将 @torch.jit.export 装饰器添加到方法上。要选择不编译，请使用 @torch.jit.ignore 或 @torch.jit.unused。

示例（模块中导出和忽略的方法）：

import torch
import torch.nn as nnclass MyModule(nn.Module):def __init__(self):super().__init__()@torch.jit.exportdef some_entry_point(self, input):return input + 10@torch.jit.ignoredef python_only_fn(self, input):# 此函数不会被编译，因此可以使用任何 Python APIimport pdbpdb.set_trace()def forward(self, input):if self.training:self.python_only_fn(input)return input * 99scripted_module = torch.jit.script(MyModule())
print(scripted_module.some_entry_point(torch.randn(2, 2)))
print(scripted_module(torch.randn(2, 2)))

示例（使用示例输入注释 nn.Module 的 forward 方法）：

import torch
import torch.nn as nn
from typing import NamedTupleclass MyModule(NamedTuple):result: List[int]class TestNNModule(torch.nn.Module):def forward(self, a) -> MyModule:result = MyModule(result=a)return resultpdt_model = TestNNModule()# 在提供的输入下运行 pdt_model 并注释 forward 的参数
scripted_model = torch.jit.script(pdt_model, example_inputs={pdt_model: [([10, 20, ], ), ], })# 使用实际输入运行 scripted_model
print(scripted_model([20]))