I am trying to understand how to use checkpoints to optimize my training. My basic understanding was that it trades increased compute for a lower memory footprint (by re-computing instead of storing data for the backward pass). Naively then I would assume that any time I use it I should decrease memory use and increase compute time. As a first pass I plugged it into a model I am training which uses a pre-trained resnet with an additional trainable parameter after it (I am using the CPC InfoNCE loss to pretrain a vision model). I noticed that it not only produced a significant decrease in memory use but also a substantial decrease in training time. I was able to increase batch size 8x and also decreased training time by 3x. I don’t understand why this would be, and am hoping someone can help me understand what is going on here.
Below is a self contained script that will reproduce this behavior on my machine (sorry it could be more minimal, but I figure the answer will be something I don’t understand about checkpoints, not in the details of my code   ):
when I set CHECKPOINT = False it uses about 8.5G of GPU memory and runs ~2.5 batches per second. When CHECKPOINT = True it uses 1.5G of GPU memory and runs at ~8.5 batches per second…
import torch
from torch import nn
from torch.utils.data import Dataset, DataLoader
import numpy as np
from tqdm.notebook import tqdm

from torch import optim
import torchvision.models as models
from torch import nn

CHECKPOINT = True
BATCH_SIZE = 32
dev = "cuda:0"

class ImageDataset(Dataset):
    def __init__(self,length = 100000,size = 244):
        self.length = length
        self.size = 244
    def __len__(self):
        return self.length
    def __getitem__(self,idx,display = False):
        return torch.from_numpy(np.random.randn(2,3,self.size,self.size))
train = ImageDataset()
trainloader = DataLoader(
    train,
    batch_size = BATCH_SIZE,
    num_workers = 24,
    pin_memory = True
)

resnet = models.resnet50(pretrained = False)

class MODEL(nn.Module):
    def __init__(self,model):
        super(MODEL,self).__init__()
        self.model = model
        self.LR = nn.Linear(1000,1000)
    def forward(self,x):
        if CHECKPOINT == False:
            o1 = self.model(x[:,0])
            o2 = self.model(x[:,1])
        else:
            o1 = torch.utils.checkpoint.checkpoint(self.model,x[:,0])
            o2 = torch.utils.checkpoint.checkpoint(self.model,x[:,1])
        
        return torch.mean((self.LR(o1)-o2)**2)
    
resnet = MODEL(resnet).to(dev)

optimizer = optim.SGD(resnet.parameters(),lr = .001)

for T in tqdm(trainloader):
    out = torch.mean(resnet(T.float().to(dev)))
    optimizer.zero_grad()
    out.backward()
    optimizer.step()