CudaGraphModule

class 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 可调用对象的 CUDA 图形包装器。

CudaGraphModule 是一个包装器类，为 PyTorch 可调用对象提供 CUDA 图形的友好界面。

警告

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)
  

• 该函数必须执行可以使用相同缓冲区完全重新运行的代码。这意味着不支持动态形状（输入或代码执行过程中形状发生变化）。换句话说，输入必须具有恒定的形状。

• 函数的输出必须是分离的。如果需要调用优化器，请将其放在输入函数中。例如，以下函数是有效的运算符

  >>> 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_out，以防输入是 tensordict。

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

警告

CudaGraphModule 不是 Module，这是有意为之，目的是阻止收集输入模块的参数并将它们传递给优化器。

参数:

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

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

• in_keys (NestedKeys 列表) –

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

  注意

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

• out_keys (NestedKeys 列表) – 输出键，如果模块接收和输出 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)