Rohan Sikand

← back to home, blog

A Simple functorch Example

In recent years, there has been a small movement of people trying to go from stateful (Python OOP, class-based modules) to stateless (pure functions) neural network code. The standard PyTorch nn.module is indeed a OOP-based class. But more recent libraries such as JAX, introduce the ability to feasiby create stateless (just functions!) machine learning models. Now, in PyTorch version 1.13, we have functorch in-tree (in the main package). Why stateless? Read this blog post for the differences, but some reasons for why I like stateless code is because:

Less leaky abstractions (and less unknown abstractions in general!)
Closer to the mathematical form (after all, a neural network is just a series of functions chained together!)
- When you learn SGD in class in the mathematical form and then use PyTorch, the disconnect is fairly evident.
Less compute overhead (less things to keep track of internally –> less memory needed)
Ability to work a lower level (which, in my opinion, can help facilitate new ideas)
Ability to work with function transformations such as vmap, pmap, jit, and grad (PyTorch has grad… yes I know… but applying grad to a stateless function makes much more intuitive sense than applying it to some stateful module!).

This might sound like an advertisement for JAX (which might be coming up in a future blog post!), but it is really to set the stage for functorch. functorch is a library that allows you to accomplish nearly all of the above, but in PyTorch! The basic idea is to purify stateful PyTorch modules into stateless functions like this (source):

import torch
                import functorch
                from functorch import make_functional

                model = torch.nn.Linear(3, 3)
                func_model, params = make_functional(model)

As functorch is relatively new, there aren’t many examples out there showing how to use the library. So the goal of the rest of this post is to provide a simple example for creating an image classifier using functorch and PyTorch and updating the weights using SGD (no torch.optim!).

Here is the code:

import torch
                import torchplate
                from torchplate import experiment
                from torchplate import utils
                import functorch
                from functorch import grad, grad_and_value
                import torch.nn as nn
                import torch.nn.functional as F
                from torch.utils.data import DataLoader
                from rsbox import ml
                # import torchopt
                import requests
                from tqdm.auto import tqdm
                import cloudpickle as cp
                from urllib.request import urlopen
                
                
                
                class Net(nn.Module):
                    def __init__(self):
                        super().__init__()
                        self.fc1 = nn.Linear(3*32*32, 120)
                        self.fc2 = nn.Linear(120, 84)
                        self.fc3 = nn.Linear(84, 16)
                        self.fc4 = nn.Linear(16, 3)
                
                    def forward(self, x):
                        x = torch.flatten(x, 1)
                        x = F.relu(self.fc1(x))
                        x = F.relu(self.fc2(x))
                        x = F.relu(self.fc3(x))
                        x = self.fc4(x)
                        return x
                
                
                
                class OptExp:
                    def __init__(self):
                        self.model_module = Net()
                        self.criterion = nn.CrossEntropyLoss()
                        dataset = cp.load(urlopen("https://stanford.edu/~rsikand/assets/datasets/mini_cifar.pkl"))
                        self.trainloader, self.testloader = torchplate.utils.get_xy_loaders(dataset)
                        self.model, self.params = functorch.make_functional(self.model_module)  # init network
                
                
                    def predict(self, x):
                        """returns logits"""
                        assert self.model is not None
                        assert self.params is not None
                        logits = self.model(self.params, x)
                        return logits
                
                
                    @staticmethod
                    def sgd_step(params, gradients, lr):
                        """one gradient step for updating the weights"""
                        updated_params = []
                        for param, gradient in zip(params, gradients):
                            update = param - (lr * gradient)
                            updated_params.append(update)
                
                        return tuple(updated_params)
                
                
                    @staticmethod
                    def stateless_loss(params, model, criterion, batch):
                        """
                        Need to perform forward pass and loss calculation in one function
                        since we need gradients w.r.t params (must be args[0]). The first
                        value we return also needs to be the scalar loss value.
                        """
                        x, y = batch
                        logits = model(params, x)
                        loss_val = criterion(logits, y)
                        return loss_val, logits
                
                
                
                    @staticmethod
                    def train_step(params, model, criterion, batch, lr):
                        """Combine this all into one function for modularity"""
                        # has_aux means we can return more than just the scalar loss
                        grad_and_loss_fn = grad_and_value(OptExp.stateless_loss, has_aux=True)
                        grads, aux_outputs = grad_and_loss_fn(params, model, criterion, batch)  # get the grads
                        loss_val, logits = aux_outputs
                        params = OptExp.sgd_step(params, grads, lr)
                        return params, loss_val, logits
                
                
                    def train(self, num_epochs=10, lr=0.01):
                        print('Beginning training!')
                        epoch_num = 0
                        for epoch in range(num_epochs):
                            running_loss = 0.0
                            epoch_num += 1
                            tqdm_loader = tqdm(self.trainloader)
                            for batch in tqdm_loader:
                                tqdm_loader.set_description(f"Epoch {epoch_num}")
                
                                # update params with one step
                                self.params, loss_val, logits = OptExp.train_step(self.params, self.model, self.criterion, batch, lr)
                
                                running_loss += loss_val
                
                            # print loss
                            epoch_avg_loss = running_loss/len(self.trainloader)
                            print("Training Loss (epoch " + str(epoch_num) + "):", epoch_avg_loss)
                
                
                        print('Finished training!')
                
                
                exp = OptExp()
                exp.train(num_epochs=50, lr=0.01)