Hi all,
I’m trying to train a model on my GPU (RTX 2080 super) using Gradient Checkpointing in order to significantly reduce the usage of VRAM. I’m using torch.utils.checkpoint.checkpoint.
The model in which I want to apply it is a simple CNN with a flatten layer at the end. Although I think I applied it right I’m not having any memory usage reduction. The memory usage with Gradient Checkpointing is the same as without it, however I do see a increase in the time per epoch (something expected given the nature of Gradient Checkpointing).
I’ve been searching through the forums but I didn’t see anything related with this.
The model definition is the following:
class my_model(torch.nn.Module):

    def __init__(self):
        super(my_model, self).__init__()

        # Input shape [19, 22, 20] -> Output shape [3258]
        self.conv1 = torch.nn.Conv2d(in_channels = 20,
                                     out_channels = 50,
                                     kernel_size = 3,
                                     padding = 1)

        self.conv2 = torch.nn.Conv2d(in_channels = 50,
                                     out_channels = 25,
                                     kernel_size = 3,
                                     padding = 1)

        self.conv3 = torch.nn.Conv2d(in_channels = 25,
                                     out_channels = 10,
                                     kernel_size = 3,
                                     padding = 1)

        self.fc1 = torch.nn.Linear(in_features = 19 * 22 * 10,
                                   out_features = 3258)

    def flat(self, x):
        return torch.flatten(x, 1)

    def forward(self, x):

        x = self.conv1(x)
        x = F.relu(x)

        x = checkpoint(self.conv2, x)
        x = checkpoint(F.relu, x)

        x = checkpoint(self.conv3, x)
        x = checkpoint(F.relu, x)

        x = checkpoint(self.flat, x)
        x = checkpoint(self.fc1, x)

        return x

I understand that when you checkpoint some operation the activations of that operation are not saved in memory (it’s recomputed when needed by the backward operation).
For show purposes I define some random train and validation data:
class MyDataset(Dataset):

    def __init__(self, data, target, transform = None):
        self.data = torch.from_numpy(data).float()
        self.target = torch.from_numpy(target).float()
        self.transform = transform

    def __getitem__(self, index):
        x = self.data[index]
        y = self.target[index]

        if self.transform:
             x = self.transform(x)

        return x, y

    def __len__(self):
        return len(self.data)

# Random data
x = np.random.randn(8766, 20, 19, 22) # 19 x 22 x 20 array
y = np.random.randn(8766, 3258) # 3258 vector

# Generate train and validation indices
validation_split = 0.1
idxs = list(range(x.shape[0]))
np.random.shuffle(idxs)

split = int(np.floor(validation_split * len(idxs)))
train_idxs, val_idxs = idxs[split:], idxs[:split]

Train and Validation DataLoaders:
# Train dataloader
dataset_train = MyDataset(data = x[train_idxs, :, :, :],
                          target = y[train_idxs, :])
loader_train = DataLoader(dataset_train,
                          batch_size = 100,
                          shuffle = True,
                          num_workers = 1)

# Validation dataloader
dataset_val = MyDataset(data = x[val_idxs, :, :, :],
                        target = y[val_idxs, :])
loader_val = DataLoader(dataset_val,
                        batch_size = 100,
                        shuffle = True,
                        num_workers = 1)

And I finally I train the model:
# Instanciate the model
model = my_model.cuda()

# Initialize lowest loss value registered (validation)
lowest_loss_val = 10000

# Initialize the number of epochs without decreasement in loss (validation)
epochs_loss_val = 1

# List for the per epoch losses in train
avg_train_loss_list = []

# List for the per epoch losses in validation
avg_val_loss_list = []

# Model hyperparameters
epochs = 100000
learning_rate = 0.0001
patience = 30 # Early stopping patience

# Loss function and optimizer
loss_function = torch.nn.MSELoss()
optimizer = optim.Adam(model.parameters(),
                       lr = learning_rate)

# Iterate over epochs
for epoch in range(epochs):

    # Record start time of epoch
    x1 = time.time()

    # Training step
    model.train()
    loss_train_list = []

    for batch_idx, (data, target) in enumerate(loader_train):

        # Clean gradients value
        optimizer.zero_grad()

        # Forward pass
        output = model(data.cuda())

        # Loss calculation
        loss = loss_function(output, target.cuda())

        # Append batch loss
        loss_train_list.append(loss.item())

        # Gradients calculation
        loss.backward()

        # Optimizer update
        optimizer.step()

    # Calculate total time spent on epoch
    x2 = time.time() - x1

    # Training loss averaged over batches (loss per epoch)
    avg_train_loss_list.append(np.average(loss_train_list))

    # Validation step (earlystopping purposes)
    model.eval()
    loss_val_list = []

    for batch_idx, (data, target) in enumerate(loader_val):

        # Forward pass
        output = model(data.cuda())

        # Loss calculation
        loss = loss_function(output, target.cuda())

        # Append batch loss
        loss_val_list.append(loss.item())

    # Validation loss averaged over batches
    avg_val_loss_list.append(np.average(loss_val_list))

    # Log info
    print('Epoch {} | Training Loss {} | Validation Loss {} | Elapsed seconds {}'.format(epoch + 1,
                                                                    round(avg_train_loss_list[-1], 5),
                                                                    round(avg_val_loss_list[-1], 5),
                                                                    round(x2, 5)))

    # Earlystopping checking
    # If latest loss is lowest than the lower one
    if avg_val_loss_list[-1] < lowest_loss_val:

        lowest_loss_val = avg_val_loss_list[-1]
        epochs_loss_val = 1

        # Save model
        torch.save(model.state_dict(), 'model.pt')

    # If we've been without losses decreasement for patience epochs
    elif epochs_loss_val > patience:

        print('{} epochs without loss reduction. Finishing training...'.format(epochs_loss_val - 1))
        print('Final loss: {}'.format(round(lowest_loss_val, 5)))
        break # Break training loop

    epochs_loss_val += 1

I’ve been trying to find without success the reason why this is not behaving as expected. Is my implementation wrong? Can’t my model’s architecture benefit from Gradient Checkpointing? I would appreciate any help.
I’m running the code on Pytorch 1.31
Thank you in advance!!