functorch.make_functional

functorch.make_functional(model) → func, params

Given a torch.nn.Module, make_functional() extracts the state (params) and returns a functional version of the model, func. This makes it so that it is possible use transforms over the parameters of model.

func can be invoked as follows:

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)

And here is an example of applying the grad transform over the parameters of a model.

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)

If the model has any buffers, please use make_functional_with_buffers() instead.