torch.jit.fork¶
- torch.jit.fork(func, *args, **kwargs)[source]¶
创建执行 func 的异步任务以及对该执行结果值的引用。
fork 将立即返回,因此 func 的返回值可能尚未计算。要强制完成任务并访问返回值,请在 Future 上调用 torch.jit.wait。使用返回 T 的 func 调用的 fork 的类型为 torch.jit.Future[T]。可以任意嵌套 fork 调用,并且可以使用位置参数和关键字参数调用。仅当在 TorchScript 中运行时才会发生异步执行。如果在纯 Python 中运行,fork 不会并行执行。当在跟踪期间调用时,fork 也不会并行执行,但是 fork 和 wait 调用将被捕获到导出的 IR 图中。
警告
fork 任务将以非确定性方式执行。我们建议仅为不修改其输入、模块属性或全局状态的纯函数生成并行的 fork 任务。
- 参数
func (callable 或 torch.nn.Module) – 将被调用的 Python 函数或 torch.nn.Module。如果在 TorchScript 中执行,它将异步执行,否则不会。跟踪的 fork 调用将被捕获到 IR 中。
*args – 用来调用 func 的参数。
**kwargs – 用来调用 func 的参数。
- 返回值
对 func 执行的引用。只有通过 torch.jit.wait 强制完成 func 才能访问值 T。
- 返回类型
torch.jit.Future[T]
示例(fork 一个自由函数)
import torch from torch import Tensor def foo(a : Tensor, b : int) -> Tensor: return a + b def bar(a): fut : torch.jit.Future[Tensor] = torch.jit.fork(foo, a, b=2) return torch.jit.wait(fut) script_bar = torch.jit.script(bar) input = torch.tensor(2) # only the scripted version executes asynchronously assert script_bar(input) == bar(input) # trace is not run asynchronously, but fork is captured in IR graph = torch.jit.trace(bar, (input,)).graph assert "fork" in str(graph)
示例(fork 一个模块方法)
import torch from torch import Tensor class AddMod(torch.nn.Module): def forward(self, a: Tensor, b : int): return a + b class Mod(torch.nn.Module): def __init__(self) -> None: super(self).__init__() self.mod = AddMod() def forward(self, input): fut = torch.jit.fork(self.mod, a, b=2) return torch.jit.wait(fut) input = torch.tensor(2) mod = Mod() assert mod(input) == torch.jit.script(mod).forward(input)