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

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 代码,并返回一个 ScriptModuleScriptFunction。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

相关推荐
孙同学要努力2 小时前
全连接神经网络案例——手写数字识别
人工智能·深度学习·神经网络
Eric.Lee20212 小时前
yolo v5 开源项目
人工智能·yolo·目标检测·计算机视觉
其实吧33 小时前
基于Matlab的图像融合研究设计
人工智能·计算机视觉·matlab
丕羽3 小时前
【Pytorch】基本语法
人工智能·pytorch·python
ctrey_3 小时前
2024-11-1 学习人工智能的Day20 openCV(2)
人工智能·opencv·学习
bryant_meng3 小时前
【python】Distribution
开发语言·python·分布函数·常用分布
SongYuLong的博客3 小时前
Air780E基于LuatOS编程开发
人工智能
Jina AI3 小时前
RAG 系统的分块难题:小型语言模型如何找到最佳断点?
人工智能·语言模型·自然语言处理
-派神-3 小时前
大语言模型(LLM)量化基础知识(一)
人工智能·语言模型·自然语言处理
johnny_hhh3 小时前
AI大模型重塑软件开发流程:定义、应用场景、优势、挑战及未来展望
人工智能