torch.cuda.make_graphed_callables

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

接受可调用对象（函数或 nn.Module）并返回图化版本。

每个图化可调用对象的前向传播将其源可调用对象的前向 CUDA 工作作为 CUDA 图在一个 autograd 节点内运行。

图化可调用对象的前向传播还会向 autograd 图添加一个反向传播节点。在反向传播期间，该节点将其可调用对象的反向传播工作作为 CUDA 图运行。

因此，在启用 autograd 的训练循环中，每个图化可调用对象都应是其源可调用对象的直接替代品。

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

如果您传递多个可调用对象的元组，它们的捕获将使用相同的内存池。有关何时适用此情况，请参阅 图内存管理。

参数

• callables (torch.nn.Module 或 Python 函数，或 元组 of these) – 要图化（graph）的可调用对象或可调用对象列表。有关何时适合传递可调用对象的元组，请参阅 图内存管理。如果您传递可调用对象的元组，它们在元组中的顺序必须与它们在实际工作负载中运行的顺序相同。

• sample_args (Tensor 的 元组，或 Tensor 的 元组 的 元组) – 每个可调用对象的示例参数。如果传递了单个可调用对象，则 sample_args 必须是参数 Tensor 的单个元组。如果传递了可调用对象的元组，则 sample_args 必须是参数 Tensor 的元组的元组。

• num_warmup_iters (int) – 热身迭代的次数。当前，DataDistributedParallel 需要 11 次迭代进行热身。默认值：3。

• allow_unused_input (bool) – 如果为 False，则指定在计算输出时未使用（因此其梯度始终为零）的输入将导致错误。默认值为 False。

• pool (可选) – 令牌（由 graph_pool_handle() 或 other_Graph_instance.pool() 返回），提示此图可能与指定的内存池共享内存。请参阅 图内存管理。

注意

sample_args 中每个 Tensor 的 requires_grad 状态必须与训练循环中相应的实际输入所期望的状态一致。

警告

此 API 处于 beta 阶段，未来版本中可能会发生变化。

警告

每个可调用对象的 sample_args 必须只包含 Tensor。不允许其他类型。

警告

返回的可调用对象不支持高阶微分（例如，二次反向传播）。

警告

传递给 make_graphed_callables() 的任何 Module 中，只有参数可以是可训练的。Buffer 必须设置 requires_grad=False。

警告

将 torch.nn.Module 传递给 make_graphed_callables() 后，您不能添加或移除该 Module 的任何参数或 buffer。

警告

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

警告

运行图化可调用对象时，您必须按照其 sample_args 中参数出现的顺序和格式传递参数。

警告

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