CudaGraphModule

类 tensordict.nn.CudaGraphModule(module: Callable[[Union[List[Tensor], TensorDictBase]], None], warmup: int = 2, in_keys: Optional[List[NestedKey]] = None, out_keys: Optional[List[NestedKey]] = None)

一个用于 PyTorch 可调用对象的 cudagraph 包装器。

CudaGraphModule 是一个包装类，它为 PyTorch 可调用对象提供了用于 CUDA graphs 的用户友好接口。

警告

CudaGraphModule 是一个原型功能，其 API 限制未来可能会发生变化。

此类提供了用户友好的 cudagraphs 接口，允许在 GPU 上快速执行操作，且没有 CPU 开销。它对函数的输入进行基本检查，并提供类似于 nn.Module 的 API 来运行

警告

此模块要求被包装的函数满足一些要求。用户有责任确保所有这些要求都得到满足。

• 该函数不能有动态控制流。例如，以下代码片段将无法被包装在 CudaGraphModule 中

  >>> def func(x):
  ...     if x.norm() > 1:
  ...         return x + 1
  ...     else:
  ...         return x - 1
  

  幸运的是，PyTorch 在大多数情况下都提供了解决方案

  >>> def func(x):
  ...     return torch.where(x.norm() > 1, x + 1, x - 1)
  

• 该函数必须执行可以使用相同缓冲区精确重新运行的代码。这意味着不支持动态形状（输入或代码执行期间形状发生变化）。换句话说，输入必须具有常量形状。

• 函数的输出必须是 detached（分离）的。如果需要调用优化器，请将其放在输入函数中。例如，以下函数是一个有效的操作符

  >>> def func(x, y):
  ...     optim.zero_grad()
  ...     loss_val = loss_fn(x, y)
  ...     loss_val.backward()
  ...     optim.step()
  ...     return loss_val.detach()
  

• 输入不应该可微分。如果您需要使用 nn.Parameters（或通常是可微分的张量），只需编写一个函数，将它们用作全局值，而不是作为输入传递

  >>> x = nn.Parameter(torch.randn(()))
  >>> optim = Adam([x], lr=1)
  >>> def func(): # right
  ...     optim.zero_grad()
  ...     (x+1).backward()
  ...     optim.step()
  >>> def func(x): # wrong
  ...     optim.zero_grad()
  ...     (x+1).backward()
  ...     optim.step()
  

• 作为张量或 tensordict 的 args 和 kwargs 可以改变（前提是设备和形状匹配），但非张量的 args 和 kwargs 不应改变。例如，如果函数接收一个字符串输入，并且输入在任何时候发生变化，模块将静默地使用捕获 cudagraph 时使用的字符串执行代码。唯一支持的关键字参数是输入为 tensordict 时的 tensordict_out。

• 如果模块是 TensorDictModuleBase 实例，并且输出 id 与输入 id 匹配，则在调用 CudaGraphModule 期间会保留此 identity。在所有其他情况下，无论其元素是否与其中一个输入匹配，输出都将被克隆。

警告

CudaGraphModule 的设计不是 Module，以避免收集输入模块的参数并将其传递给优化器。

参数：

• module (可调用对象) – 一个函数，接收张量（或 tensordict）作为输入，并输出一个 PyTreeable 的张量集合。如果提供了 tensordict，模块也可以使用关键字参数运行（参见下面的示例）。

• warmup (int, 可选) – 模块被编译时的 warmup 步骤数（编译后的模块在被 cudagraphs 捕获之前应该运行几次）。所有模块的默认值为 2。

• in_keys (NestedKeys 列表) –

  输入键，如果模块以 TensorDict 作为输入。如果 module.in_keys 存在，则默认为该值，否则为 None。

  注意

  如果提供了 in_keys 但为空，则假定模块接收 tensordict 作为输入。这足以让 CudaGraphModule 知道该函数应被视为 TensorDictModule，但关键字参数将不会被分发。请参见下面的示例。

• out_keys (NestedKeys 列表) – 输出键，如果模块以 TensorDict 作为输入并输出 TensorDict。如果 module.out_keys 存在，则默认为该值，否则为 None。

示例

>>> # Wrap a simple function
>>> def func(x):
...     return x + 1
>>> func = CudaGraphModule(func)
>>> x = torch.rand((), device='cuda')
>>> out = func(x)
>>> assert isinstance(out, torch.Tensor)
>>> assert out == x+1
>>> # Wrap a tensordict module
>>> func = TensorDictModule(lambda x: x+1, in_keys=["x"], out_keys=["y"])
>>> func = CudaGraphModule(func)
>>> # This can be called either with a TensorDict or regular keyword arguments alike
>>> y = func(x=x)
>>> td = TensorDict(x=x)
>>> td = func(td)