The models I’m trying to train doesn’t fit on my GPU memory. I’ve already implemented Gradient Accumulation and FP16 training but it’s not enough, that’s why I’m using Gradient Checkpointing. I don’t care about computation time, I just want to fit my model on memory.
I tried checkpointing the conv + relu operations:
class my_model(torch.nn.Module):

    def __init__(self):
        super(CNN10_GC, 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 block1(self, x):
        return F.relu(self.conv2(x))

    def block2(self, x):
        return F.relu(self.conv3(x))

    def block3(self, x):
        return self.fc1(self.flat(x))

    def forward(self, x):

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

        x = checkpoint(self.block1, x)

        x = checkpoint(self.block2, x)

        x = checkpoint(self.block3, x)

        return x

And even all the convs together:
class my_model(torch.nn.Module):

    def __init__(self):
        super(CNN10_GC, 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 block1(self, x):

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

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

        return x

    def forward(self, x):

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

        x = checkpoint(self.block1, x)

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

        return x

But the memory usage is exactly the same, however as I said I notice a little time increase per epoch.
I was thinking that maybe the activations not calculated when using the checkpointing function were not big enough so the memory saving was not meaningful. I checked this by setting the requires_grad attribute of the input tensor in the forward method to True but I did see a increase in memory, so my theory was wrong.
I have not idea what is failing in my models… appreciate any help!