functorch.make_functional_with_buffers

functorch.make_functional_with_buffers(model, disable_autograd_tracking=False)

make_functional(model, disable_autograd_tracking=False) -> func, params

给定一个 torch.nn.Module，make_functional() 提取状态（参数）并返回模型的功能版本 func。这使得可以使用模型 model 参数上的变换。

func 可以按如下方式调用

import torch
import torch.nn as nn
from functorch import make_functional

x = torch.randn(4, 3)
model = nn.Linear(3, 3)
func, params = make_functional(model)
func(params, x)

以下是如何在模型参数上应用 grad 变换的示例。

import torch
import torch.nn as nn
from functorch import make_functional, grad

x = torch.randn(4, 3)
t = torch.randn(4, 3)
model = nn.Linear(3, 3)
func, params = make_functional(model)

def compute_loss(params, x, t):
    y = func(params, x)
    return nn.functional.mse_loss(y, t)

grad_weights = grad(compute_loss)(params, x, t)

如果模型有任何缓冲区，请改用 make_functional_with_buffers()。

参数

• model (torch.nn.Module) – 输入模型。

• disable_autograd_tracking (bool) – 用于禁用输出参数的梯度跟踪的标志。返回的参数与原始模型中的参数集无关。如果为 False（默认），则参数将具有 requires_grad=True（即它们可以通过常规 PyTorch 自动梯度进行跟踪），与原始模型参数的 requires_grad-ness 相匹配。否则，返回的参数将具有 requires_grad=False。默认值为 False。如果您计划使用常规 PyTorch 自动梯度（例如，如果您想调用 .backward() 或 torch.autograd.grad()，请将 disable_autograd_tracking=False。否则，如果您只计划使用 functorch 的梯度变换，请将 disable_autograd_tracking=True 设置为避免使用 PyTorch 自动梯度不必要地跟踪历史记录。

警告

我们将 functorch 集成到了 PyTorch 中。作为集成的最后一步，functorch.make_functional_with_buffers 已在 PyTorch 2.0 中弃用，将在 PyTorch >= 2.3 的未来版本中删除。请改用 torch.func.functional_call；有关更多详细信息，请参阅 PyTorch 2.0 发行说明和/或 torch.func 迁移指南 https://pytorch.ac.cn/docs/master/func.migrating.html