torch.jit.script¶
- torch.jit.script(obj, optimize=None, _frames_up=0, _rcb=None, example_inputs=None)[source][source]¶
脚本化函数。
脚本化函数或
nn.Module
将检查源代码,使用 TorchScript 编译器将其编译为 TorchScript 代码,并返回ScriptModule
或ScriptFunction
。TorchScript 本身是 Python 语言的子集,因此并非 Python 中的所有功能都有效,但我们提供了足够的功能来计算张量并执行控制依赖操作。有关完整指南,请参阅 TorchScript 语言参考。脚本化字典或列表会将其中包含的数据复制到 TorchScript 实例中,然后可以在 Python 和 TorchScript 之间通过引用传递,且零复制开销。
torch.jit.script
可以用作模块、函数、字典和列表的函数以及用作
@torch.jit.script
装饰器,用于 TorchScript 类 和函数。
- 参数
obj (Callable, class, 或 nn.Module) – 要编译的
nn.Module
、函数、类类型、字典或列表。example_inputs (Union[List[Tuple], Dict[Callable, List[Tuple]], None]) – 提供示例输入以注释函数或
nn.Module
的参数。
- 返回值
如果
obj
是nn.Module
,则script
返回ScriptModule
对象。返回的ScriptModule
将具有与原始nn.Module
相同的子模块和参数集。如果obj
是独立函数,则将返回ScriptFunction
。如果obj
是dict
,则script
返回 torch._C.ScriptDict 的实例。如果obj
是list
,则script
返回 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 # See the compiled graph as Python code print(foo.code) # Call the function using the TorchScript interpreter foo(torch.ones(2, 2), torch.ones(2, 2))
- **使用 example_inputs 脚本化函数
示例输入可用于注释函数参数。
示例 (在脚本化之前注释函数)
import torch def test_sum(a, b): return a + b # Annotate the arguments to be int scripted_fn = torch.jit.script(test_sum, example_inputs=[(3, 4)]) print(type(scripted_fn)) # torch.jit.ScriptFunction # See the compiled graph as Python code print(scripted_fn.code) # Call the function using the TorchScript interpreter scripted_fn(20, 100)
- 脚本化 nn.Module
默认情况下,脚本化
nn.Module
将编译forward
方法,并递归编译forward
调用的任何方法、子模块和函数。如果nn.Module
仅使用 TorchScript 中支持的功能,则无需更改原始模块代码。script
将构造ScriptModule
,该模块具有原始模块的属性、参数和方法的副本。示例 (脚本化带有 Parameter 的简单模块)
import torch class MyModule(torch.nn.Module): def __init__(self, N, M): super().__init__() # This parameter will be copied to the new ScriptModule self.weight = torch.nn.Parameter(torch.rand(N, M)) # When this submodule is used, it will be compiled self.linear = torch.nn.Linear(N, M) def forward(self, input): output = self.weight.mv(input) # This calls the `forward` method of the `nn.Linear` module, which will # cause the `self.linear` submodule to be compiled to a `ScriptModule` here output = self.linear(output) return output scripted_module = torch.jit.script(MyModule(2, 3))
示例 (脚本化带有跟踪子模块的模块)
import torch import torch.nn as nn import torch.nn.functional as F class MyModule(nn.Module): def __init__(self) -> None: super().__init__() # torch.jit.trace produces a ScriptModule's conv1 and 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) -> None: super().__init__() @torch.jit.export def some_entry_point(self, input): return input + 10 @torch.jit.ignore def python_only_fn(self, input): # This function won't be compiled, so any # Python APIs can be used 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)))
示例 (使用 example_inputs 注释 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() # Runs the pdt_model in eager model with the inputs provided and annotates the arguments of forward scripted_model = torch.jit.script(pdt_model, example_inputs={pdt_model: [([10, 20, ], ), ], }) # Run the scripted_model with actual inputs print(scripted_model([20]))