本文主要是介绍【官方文档解读】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]))
官方文档链接:https://pytorch.org/docs/stable/generated/torch.jit.script.html#torch.jit.script
这篇关于【官方文档解读】torch.jit.script 的使用,并附上官方文档中的示例代码的文章就介绍到这儿,希望我们推荐的文章对编程师们有所帮助!