由 OpenMMLab 的部署教程 所述，对于模型中存在有控制条件的（如 if，for 等），需要用 torch.jit.script 而非采样默认的 torch.jit.trace 方法。本文则详细介绍了下官方文档中对 torch.jit.script 的解释和示例代码。
torch.jit.script
torch.jit.script 用于将函数或 nn.Module 编译为 TorchScript。
函数签名
python
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。
示例（脚本化函数）：
python
import torch
@torch.jit.script
def foo(x, y):
if x.max() > y.max():
r = x
else:
r = y
return r
print(type(foo)) # torch.jit.ScriptFunction
# 以 Python 代码查看编译后的图
print(foo.code)
# 使用 TorchScript 解释器调用函数
foo(torch.ones(2, 2), torch.ones(2, 2))使用示例输入脚本化函数
示例输入可用于注释函数参数。
示例（脚本化前注释函数）：
python
import torch
def 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，其中包含原始模块的属性、副本和方法。
示例（脚本化包含参数的简单模块）：
python
import torch
class MyModule(torch.nn.Module):
def __init__(self, N, M):
super().__init__()
# 此参数将被复制到新的 ScriptModule
self.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 output
scripted_module = torch.jit.script(MyModule(2, 3))示例（脚本化包含 traced 子模块的模块）：
python
import torch
import torch.nn as nn
import torch.nn.functional as F
class MyModule(nn.Module):
def __init__(self):
super().__init__()
# torch.jit.trace 生成一个 ScriptModule 的 conv1 和 conv2
self.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 input
scripted_module = torch.jit.script(MyModule())要编译 forward 以外的方法（并递归编译它调用的任何内容），请将 @torch.jit.export 装饰器添加到方法上。要选择不编译，请使用 @torch.jit.ignore 或 @torch.jit.unused。
示例（模块中导出和忽略的方法）：
python
import torch
import torch.nn as nn
class MyModule(nn.Module):
def __init__(self):
super().__init__()
@torch.jit.export
def some_entry_point(self, input):
return input + 10
@torch.jit.ignore
def python_only_fn(self, input):
# 此函数不会被编译，因此可以使用任何 Python API
import pdb
pdb.set_trace()
def forward(self, input):
if self.training:
self.python_only_fn(input)
return input * 99
scripted_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 方法）：
python
import torch
import torch.nn as nn
from typing import NamedTuple
class MyModule(NamedTuple):
result: List[int]
class TestNNModule(torch.nn.Module):
def forward(self, a) -> MyModule:
result = MyModule(result=a)
return result
pdt_model = TestNNModule()
# 在提供的输入下运行 pdt_model 并注释 forward 的参数
scripted_model = torch.jit.script(pdt_model, example_inputs={pdt_model: [([10, 20, ], ), ], })
# 使用实际输入运行 scripted_model
print(scripted_model([20]))官方文档链接：https://pytorch.org/docs/stable/generated/torch.jit.script.html#torch.jit.script