【官方文档解读】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 = 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))

使用示例输入脚本化函数

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

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

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，其中包含原始模块的属性、副本和方法。

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

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 子模块的模块）：

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。

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

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 方法）：

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