TorchScript 中的动态并行性¶

创建时间： 2020年7月28日 |上次更新时间：2024 年 12 月 2 日 |上次验证： Nov 05， 2024

警告

TorchScript 不再处于积极开发阶段。

在本教程中，我们将介绍在 TorchScript 中执行动态互操作并行性的语法。此并行度具有以下属性：

dynamic - 创建的并行任务的数量及其工作负载可能取决于程序的控制流。
inter-op - 并行性与并行运行 TorchScript 程序片段有关。这与 intra-op parallelism 不同，后者关注拆分单个 Operator 并并行运行 Operator 工作的子集。

基本语法¶

动态并行的两个重要 API 是：

torch.jit.fork(fn : Callable[..., T], *args, **kwargs) -> torch.jit.Future[T]
torch.jit.wait(fut : torch.jit.Future[T]) -> T

通过示例来演示这些工作原理的一个好方法是：

import torch

def foo(x):
    return torch.neg(x)

@torch.jit.script
def example(x):
    # Call `foo` using parallelism:
    # First, we "fork" off a task. This task will run `foo` with argument `x`
    future = torch.jit.fork(foo, x)

    # Call `foo` normally
    x_normal = foo(x)

    # Second, we "wait" on the task. Since the task may be running in
    # parallel, we have to "wait" for its result to become available.
    # Notice that by having lines of code between the "fork()" and "wait()"
    # call for a given Future, we can overlap computations so that they
    # run in parallel.
    x_parallel = torch.jit.wait(future)

    return x_normal, x_parallel

print(example(torch.ones(1))) # (-1., -1.)

fork()将可调用对象和参数转换为该可调用对象，并创建一个异步任务来执行。可以是 function、method 或 Module 实例。返回引用此执行的结果的值，称为 . 因为在创建异步任务后立即返回，所以可能会在调用被执行。因此，用于等待异步任务完成并返回值。fnargskwargsfnfnfork()Futureforkfnfork()wait()

这些构造可用于重叠函数（如 Worked Example 部分所示）或与其他语言组合像 Loops 这样的结构：

import torch
from typing import List

def foo(x):
    return torch.neg(x)

@torch.jit.script
def example(x):
    futures : List[torch.jit.Future[torch.Tensor]] = []
    for _ in range(100):
        futures.append(torch.jit.fork(foo, x))

    results = []
    for future in futures:
        results.append(torch.jit.wait(future))

    return torch.sum(torch.stack(results))

print(example(torch.ones([])))

注意

当我们初始化一个空的 Future 列表时，我们需要添加一个显式的键入注释。在 TorchScript 中，空容器默认为假设它们包含 Tensor 值，因此我们对列表构造函数进行注释 # 指定为futuresList[torch.jit.Future[torch.Tensor]]

此示例用于启动 100 个函数，等待 100 个任务完成，然后对结果求和，返回 .fork()foo-100.0

应用示例：双向 LSTM 的系综¶

让我们尝试将并行性应用于一个更现实的示例，看看是哪种我们可以摆脱它。首先，我们定义基线模型：一个双向 LSTM 层的集合。

import torch, time

# In RNN parlance, the dimensions we care about are:
# # of time-steps (T)
# Batch size (B)
# Hidden size/number of "channels" (C)
T, B, C = 50, 50, 1024

# A module that defines a single "bidirectional LSTM". This is simply two
# LSTMs applied to the same sequence, but one in reverse
class BidirectionalRecurrentLSTM(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.cell_f = torch.nn.LSTM(input_size=C, hidden_size=C)
        self.cell_b = torch.nn.LSTM(input_size=C, hidden_size=C)

    def forward(self, x : torch.Tensor) -> torch.Tensor:
        # Forward layer
        output_f, _ = self.cell_f(x)

        # Backward layer. Flip input in the time dimension (dim 0), apply the
        # layer, then flip the outputs in the time dimension
        x_rev = torch.flip(x, dims=[0])
        output_b, _ = self.cell_b(torch.flip(x, dims=[0]))
        output_b_rev = torch.flip(output_b, dims=[0])

        return torch.cat((output_f, output_b_rev), dim=2)


# An "ensemble" of `BidirectionalRecurrentLSTM` modules. The modules in the
# ensemble are run one-by-one on the same input then their results are
# stacked and summed together, returning the combined result.
class LSTMEnsemble(torch.nn.Module):
    def __init__(self, n_models):
        super().__init__()
        self.n_models = n_models
        self.models = torch.nn.ModuleList([
            BidirectionalRecurrentLSTM() for _ in range(self.n_models)])

    def forward(self, x : torch.Tensor) -> torch.Tensor:
        results = []
        for model in self.models:
            results.append(model(x))
        return torch.stack(results).sum(dim=0)

# For a head-to-head comparison to what we're going to do with fork/wait, let's
# instantiate the model and compile it with TorchScript
ens = torch.jit.script(LSTMEnsemble(n_models=4))

# Normally you would pull this input out of an embedding table, but for the
# purpose of this demo let's just use random data.
x = torch.rand(T, B, C)

# Let's run the model once to warm up things like the memory allocator
ens(x)

x = torch.rand(T, B, C)

# Let's see how fast it runs!
s = time.time()
ens(x)
print('Inference took', time.time() - s, ' seconds')

在我的计算机上，此网络在几秒钟内运行。我们可以做得更好！2.05

并行化前向层和后向层¶

我们可以做的一件非常简单的事情是并行化前向层和后向层在。为此，计算的结构是静态的，因此我们实际上甚至不需要任何 Loop。让我们重写这样的方法：BidirectionalRecurrentLSTMforwardBidirectionalRecurrentLSTM

def forward(self, x : torch.Tensor) -> torch.Tensor:
    # Forward layer - fork() so this can run in parallel to the backward
    # layer
    future_f = torch.jit.fork(self.cell_f, x)

    # Backward layer. Flip input in the time dimension (dim 0), apply the
    # layer, then flip the outputs in the time dimension
    x_rev = torch.flip(x, dims=[0])
    output_b, _ = self.cell_b(torch.flip(x, dims=[0]))
    output_b_rev = torch.flip(output_b, dims=[0])

    # Retrieve the output from the forward layer. Note this needs to happen
    # *after* the stuff we want to parallelize with
    output_f, _ = torch.jit.wait(future_f)

    return torch.cat((output_f, output_b_rev), dim=2)

在此示例中，将的执行委托给另一个线程当它继续执行 .这会导致执行单元格彼此重叠。forward()cell_fcell_b

使用这个简单的修改再次运行脚本会产生几秒钟的运行时间，从而将！1.7117%

旁白：可视化并行度¶

我们还没有完成对模型的优化，但值得介绍一下我们用于可视化性能。一个重要的工具是 PyTorch 分析器。

让我们使用分析器和 Chrome 跟踪导出功能来可视化并行化模型的性能：

with torch.autograd.profiler.profile() as prof:
    ens(x)
prof.export_chrome_trace('parallel.json')

此代码段将写出一个名为 .如果你将 Google Chrome 导航到，单击按钮，然后 load 到该 JSON 文件中，您应该会看到如下所示的时间线：parallel.jsonchrome://tracingLoad

时间轴的横轴表示时间，纵轴表示时间表示执行线程。正如我们所看到的，我们一次运行两个实例。这是我们努力并行化双向层！lstm

在 Ensemble 中并行化模型¶

您可能已经注意到，我们的代码：我们也可以并行运行彼此。做到这一点的方法很简单，这就是我们应该改变的方式方法：LSTMEnsembleforwardLSTMEnsemble

def forward(self, x : torch.Tensor) -> torch.Tensor:
    # Launch tasks for each model
    futures : List[torch.jit.Future[torch.Tensor]] = []
    for model in self.models:
        futures.append(torch.jit.fork(model, x))

    # Collect the results from the launched tasks
    results : List[torch.Tensor] = []
    for future in futures:
        results.append(torch.jit.wait(future))

    return torch.stack(results).sum(dim=0)

或者，如果你重视简洁，我们可以使用列表推导式：

def forward(self, x : torch.Tensor) -> torch.Tensor:
    futures = [torch.jit.fork(model, x) for model in self.models]
    results = [torch.jit.wait(fut) for fut in futures]
    return torch.stack(results).sum(dim=0)

就像 intro 中描述的，我们使用了 loop 来分叉每个模型。然后，我们使用另一个循环来等待所有要完成的任务。这提供了更多的计算重叠。

通过这个小的更新，脚本可以在几秒钟内运行，从而大大加快速度之！对于两行代码来说相当不错。1.432%

我们还可以再次使用 Chrome 跟踪器来查看发生了什么：

我们现在可以看到所有实例都完全并行运行。LSTM

结论¶

在本教程中，我们了解了和的基本 API 用于在 TorchScript 中执行动态的互操作并行性。我们看到了一些典型的使用这些函数并行执行函数、方法或 TorchScript 代码中。最后，我们解决了使用此技术优化模型并探索性能的示例 PyTorch 中提供的测量和可视化工具。fork()wait()Modules