Excessive Memory Usage

JAllsopp: <p>Hi,<br>
I’m running a CNN model on Pytorch, and I’m getting an out of memory crash after 24 batches of a training and validation stage. My batches contain 32 images (300x300 greyscale) and I’m using float32. The classes are a small json fragment. I’m using a Quadro 6000 card with 24Gb ram. The model contains 17.8 million parameters. Assuming four bytes, this shouldn’t take more than 100Mb  without gradients. If there’s a gradient per parameter, I would assume this to double, but its still well within 24Gb.</p>
<p>I think there’s probably a memory leak due to me not releasing somewhere in the training loop; Here’s the code. It fails during the validation stage.</p>
<pre><code class="lang-auto">    model.to(device=device_in_use)
    count_parameters(model)
    for epoch in range(1, n_epochs + 1):
        loss_train=0
        batch_number = 1
        for batch in train_loader:
            print(f"Batch {batch_number}")
            batch_number+=1
            images = batch[0]
            classes = batch[1]
            print(f"{images.shape=} {classes['number_of_ions'].to(torch.float32)=}")
            imgs=images.to(device=device_in_use, dtype=torch.float32)
            outputs = model(imgs)
            outputs = outputs.squeeze()
            normalised_labels = classes['number_of_ions'].to(torch.float32)/max_ions
            normalised_labels = normalised_labels.to(device=device_in_use, dtype=torch.float32)

            loss = loss_fn(outputs,normalised_labels)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            current_loss = loss.item()
            loss_train += current_loss

        writer.add_scalar("loss/training", loss_train/batch_size)
        nll =[]
        ol=[]

        validation_loss = 0

        for batch in validation_loader:
            images = batch[0]
            classes = batch[1]

            imgs= images.to(device=device_in_use, dtype=torch.float32)
            outputs = model(imgs)
            outputs = outputs.squeeze()
            normalised_labels = classes['number_of_ions'].to(torch.float32)/max_ions
            normalised_labels = normalised_labels.to(device=device_in_use, dtype=torch.float32)
            nll.extend(normalised_labels)
            ol.extend(outputs)
            loss = loss_fn(outputs,normalised_labels)
            validation_loss += loss.item()

        print(f"Epoch:{epoch} Loss:{loss_train/len(train_loader)} Validation:{validation_loss/len(validation_loader)} len train_loader:{len(train_loader)} len validation_loader: {len(validation_loader)}")
</code></pre>
<p>The model looks like this;</p>
<pre><code class="lang-auto">class Net(nn.Module):
    def __init__(self, image_size:int, last_level_output_filters:int):
        super().__init__()
        number_of_levels_of_max_pool = 3
        image_size = math.floor(image_size/pow(2,number_of_levels_of_max_pool))

        self.size_of_connection_between_cnn_and_linear = image_size*image_size*last_level_output_filters
        self.conv1 = nn.Conv2d(1,75, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(75,50, kernel_size=3, padding=1)
        self.conv3 = nn.Conv2d(50, last_level_output_filters, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm2d(75)
        self.bn2 = nn.BatchNorm2d(50)
        self.bn3 = nn.BatchNorm2d(25)
        self.dropout1 = nn.Dropout(0.2)
        self.dropout2 = nn.Dropout(0.3)
        self.fc1 = nn.Linear(self.size_of_connection_between_cnn_and_linear, 512)
        self.fc2 = nn.Linear(512,512)
        self.fc3 = nn.Linear(512,1)
        self.float()

    def forward(self, x):
        # formula for the output looks like                                                                                                                                                                                                                                                                                                               
        # batch_size* number_of_filters* ((floor((Image_width-kernel_Size+2*Padding)/stride)+1)^2), then divide by 4 (2x2) for the max_pool2d.                                                                                                                                                                                                            

        out = F.max_pool2d(self.bn1(torch.relu(self.conv1(x))),2)
        out = F.max_pool2d(self.bn2(self.dropout1(torch.relu(self.conv2(out)))),2)
        out = F.max_pool2d(self.bn3(torch.relu(self.conv3(out))),2)
        out = out.view(-1,self.size_of_connection_between_cnn_and_linear)
        out = torch.relu(self.fc1(out))
        out = self.dropout2(out)
        out = torch.relu(self.fc2(out))
        out = self.fc3(out)

        return out

</code></pre>
<p>If anyone has any advice on how to fix this memory leak or speed this up,  I would be very grateful.</p>

ptrblck: <p>You are appending the model outputs, which are still attached to the computation graph, to a <code>list</code> which will increase the memory usage since all intermediate activations will also be stored. Wrap the validation loop into a <code>no_grad()</code> context to avoid creating the computation graph or detach the output explicitly.</p>