functorch.make_functional

functorch.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 属性匹配。否则，返回的参数将具有 requires_grad=False。默认值为 False。如果您计划使用常规 PyTorch 自动微分（例如，如果您想调用 .backward() 或 torch.autograd.grad()），则将 disable_autograd_tracking 设置为 False。否则，如果您只打算使用 functorch 的梯度变换，则请将 disable_autograd_tracking 设置为 True，以避免不必要地使用 PyTorch 自动微分跟踪历史记录。

警告

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