Hi,
My GPU (RTX 2080 Super) has 8 GB of memory. Without applying any memory optimization technique it uses 1317 MiB, with Gradient Accumulation (batch size of 100 with batches of 1 element for the accumulation) uses 1097 MB and with FP16 training (using half() method) uses 987 MB. There is no decrease with Gradient Checkpointing.
You can see how this model fit on memory however a variation of it does not (the one I’m really interested in fitting on memory). This variation follows a similar architecture:
class my_model_full(torch.nn.Module):

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

        # Input shape [45, 43, 20]
        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.param1 = torch.nn.Linear(in_features = 45 * 43 * 10,
                                      out_features = 29913)

        self.param2 = torch.nn.Linear(in_features = 45 * 43 * 10,
                                      out_features = 29913)

        self.param3 = torch.nn.Linear(in_features = 45 * 43 * 10,
                                      out_features = 29913)

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

    def cat_torch(self, x1, x2, x3):
        return torch.cat((x1, x2, x3), dim = 1)

    def forward(self, x):

        x.requires_grad = True

        x = checkpoint(self.conv1, x)
        x = checkpoint(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)

        x1 = checkpoint(self.param1, x)
        x1 = checkpoint(torch.sigmoid, x1)

        x2 = checkpoint(self.param2, x)

        x3 = checkpoint(self.param3, x)

        x = checkpoint(self.cat_torch, x1, x2, x3)

        return x

Casually today I had access to a Tesla V100 with 32 GB of memory. Without any memory optimization technique I receive a OOM error, same as with Gradient Checkpointing. With Gradient Accumulation uses 32027 MB and with FP16 training uses 16943 MB. As you can see I finally achieved to fit the model on memory with Gradient Accumulation (in exchange of a big increase in time per epoch) but I still dont’have any memory reduction with Gradient Checkpointing…
I’m still curious about this issue. Appreciate any help!