torch.cuda.make_graphed_callables#

torch.cuda.make_graphed_callables(callables, sample_args, num_warmup_iters=3, allow_unused_input=False, pool=None)[source]#

接受 callables（函数或 nn.Module）并返回图版本。

每个图 callable 的前向传播会将其源 callable 的前向 CUDA 工作作为 CUDA 图在单个 autograd 节点中运行。

图 callable 的前向传播还会向 autograd 图追加一个后向节点。在后向传播期间，此节点将 callable 的后向工作作为 CUDA 图运行。

因此，每个图 callable 在启用 autograd 的训练循环中都应可作为其源 callable 的即插即用替代项。

有关详细用法和限制，请参阅部分网络捕获。

如果您传递一个包含多个 callables 的元组，它们的捕获将使用相同的内存池。有关何时适用此方法的详细信息，请参阅图内存管理。

参数

callables (torch.nn.Module 或 Python 函数，或包含这些的 元组) – 要图化的 callable 或 callables。有关传递 callables 元组的适用性，请参阅图内存管理。如果您传递 callables 元组，元组中的顺序必须与实际工作负载中的运行顺序相同。
sample_args (每个 callable 的 Tensor 元组，或 Tensor 元组的元组) – 每个 callable 的样本参数。如果您传递了单个 callable，sample_args 必须是参数 Tensor 的单个元组。如果您传递了 callables 的元组，sample_args 必须是参数 Tensor 的元组的元组。
num_warmup_iters (int) – 预热迭代次数。目前，DataDistributedParallel 需要 11 次迭代进行预热。默认为 3。
allow_unused_input (bool) – 如果为 False，则指定计算输出时未使用的输入（因此它们的 grad 始终为零）将被视为错误。默认为 False。
pool (optional) – Token（由 graph_pool_handle() 或 other_Graph_instance.pool() 返回）指示此图可能与指定池共享内存。请参阅图内存管理。

注意

sample_args 中每个 callable 的 Tensor 的 requires_grad 状态必须与其在训练循环中对应的实际输入的预期状态匹配。

警告

此 API 处于 Beta 版，未来版本中可能会更改。

警告

sample_args 对每个 callable 来说，必须只包含 Tensors。不允许其他类型。

警告

返回的 callables 不支持高阶微分（例如，二次反向传播）。

警告

在通过 make_graphed_callables() 传递的任何 Module 中，只有参数可以是可训练的。Buffer 必须具有 requires_grad=False。

警告

在通过 make_graphed_callables() 传递 torch.nn.Module 之后，您不得添加或删除该 Module 的任何参数或 buffer。

警告

在传递给 make_graphed_callables() 的 torch.nn.Module 上，在传递时不得注册模块钩子。但是，在通过 make_graphed_callables() 之后在模块上注册钩子是允许的。

警告

运行图 callable 时，您必须以与 callable 的 sample_args 中出现的顺序和格式相同的方式传递其参数。

警告

自动混合精度在 make_graphed_callables() 中仅支持禁用缓存。上下文管理器 torch.cuda.amp.autocast() 必须具有 cache_enabled=False。

文档