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 函数，或 tuple of 这些) – 要进行图优化的可调用对象或可调用对象元组。有关何时适合传递可调用对象元组，请参阅 图内存管理。如果传递可调用对象元组，则它们在元组中的顺序必须与它们在实际工作负载中运行的顺序相同。

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

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

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

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

注意

sample_args 中每个张量的 requires_grad 状态必须与训练循环中相应实际输入所期望的状态匹配。

警告

此 API 处于 beta 阶段，未来版本可能会更改。

警告

每个可调用对象的 sample_args 必须仅包含张量。不允许使用其他类型。

警告

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

警告

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

警告

在将 torch.nn.Module 传递给 make_graphed_callables() 后，您不得添加或删除该模块的任何参数或缓冲区。

警告

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

警告

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

警告

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