Working with FuncTorch: An Introduction

import wandb
import torch

# Initialize a wandb run
wandb.init(project="functorch-examples")

# Set up experiment configs
config = wandb.config
config.batch_size = 64
config.num_workers = 2
config.learning_rate = 1e-2
config.epochs = 5
config.device = "cuda:0" if torch.cuda.is_available() else "cpu"
config.classes = (
    'plane', 'car', 'bird', 'cat', 'deer',
    'dog', 'frog', 'horse', 'ship', 'truck'
)

from functorch import make_functional

model = Classifier().to(device=config.device)
# params decoupled from the model. Now functional_model is a pure function.
functional_model, params = make_functional(model)

# we don't need to store grads for backprop
# avoiding this step might result in a memory leak
for param in params:
    param.requires_grad_(False)

train_metric = Accuracy().to(config.device)
val_metric = Accuracy().to(config.device)

def compute_stateless_loss(params, inputs, labels):
    outputs = functional_model(params, inputs)
    loss = F.cross_entropy(outputs, labels)
    return loss

from functorch import grad, grad_and_value

# transformed functions
compute_gradients = grad(compute_stateless_loss)
compute_gradients_and_loss = grad_and_value(compute_stateless_loss)

# naive sgd optimizer implementation
def sgd_optimizer(weights, gradients, learning_rate):
    def step(weight, gradient):
        return weight - learning_rate * gradient

    return [step(weight, gradient) for weight, gradient in zip(weights, gradients)]

# functional train step
def functional_step(params, inputs, labels):
    # --------------------------------------------------------------------
    # compute gradients and loss using the transformed function instead of
    # outputs = model(inputs)
    # loss = criterion(outputs, labels)
    # loss.backward()
    gradients, loss = compute_gradients_and_loss(params, inputs, labels)
    # --------------------------------------------------------------------
    # update the parameters using the functional optimizer
    # instead of `optimizer.step()`
    params = sgd_optimizer(params, gradients, config.learning_rate)
    # --------------------------------------------------------------------
    return params, loss

# training function
def train(functional_step, params, metric):
    running_loss = 0.
    metric.reset()
    # look over the train dataloader
    for i, data in tqdm(enumerate(train_loader), total=len(train_loader), leave=False):
        # get the inputs; data is a list of [inputs, labels]
        inputs, labels = to_device(data, config.device)
        # perform functional train step
        params, loss = functional_step(params, inputs, labels)
        running_loss += loss.item()
        # forward propagation for updating the metric
        outputs = functional_model(params, inputs)
        metric(outputs, labels)
    return params, running_loss, metric

# validation function
def validate(params, metric):
    running_loss = 0.
    metric.reset()
    # look over the validation dataloader
    for i, data in tqdm(enumerate(test_loader), total=len(test_loader), leave=False):
        inputs, labels = to_device(data, config.device)
        outputs = functional_model(params, inputs)
        loss = F.cross_entropy(outputs, labels)
        running_loss += loss.item()
        metric(outputs, labels)
    return running_loss, metric

# loop over the dataloader multiple times
for epoch in tqdm(range(config.epochs)):
    # perform training (single loop over the train dataloader)
    params, train_loss, train_metric = train(functional_step, params, train_metric)
    # perform validation (single loop over the validation dataloader)
    val_loss, val_metric = validate(params, val_metric)
    print(f'[{epoch}] train_loss: {train_loss / 2000:.3f}, train_accuracy: {train_metric.compute():.3f}, val_loss: {val_loss / 2000:.3f}, val_accuracy: {val_metric.compute():.3f}')
    wandb.log({
        "train_loss": train_loss,
        "val_loss": val_loss,
        "train_accuracy": train_metric.compute(),
        "val_accuracy": val_metric.compute()
    })