Memory increases drasticaly when running resnet18 in CPU

username: ITs
cooked: <p>I running the following code in CPU with torch 1.7.1… My concern is that memory increases with the number of iterations. Reaches about 40GB and I can not complete not even a single epoch. Could you please explain if that is expected or how can we make it more stable?</p>
<blockquote>
<p>import torch<br>
from tqdm import tqdm<br>
import matplotlib.pyplot as plt<br>
from torchvision import transforms<br>
from torch.autograd import Variable<br>
import torchvision.datasets as pydsets<br>
from torchvision.models import resnet18</p>
<p>NUM_FEATURES = {<br>
‘cifar10’: 32*32<br>
}</p>
<p>def main():<br>
model = resnet18()<br>
print(‘Total params: %.2fM’ % (sum(p.numel() for p in model.parameters()) / 1000000.0))</p>
<pre><code>criterion = torch.nn.CrossEntropyLoss()

# optimizer = torch.optim.SGD(model.parameters(), lr=1e-1)
optimizer = torch.optim.Adam(model.parameters())

tr_losses = []
step_times = []
n_iter = 0
n_runs = 1  # For mean efficiency
total_timer.start()
for r in range(n_runs):
    for epoch in range(int(epochs)):
        running_loss = 0
        running_samples = 0
        with tqdm(total=len(train_loader.dataset)) as progressbar:
            for i, (samples, labels) in enumerate(train_loader):
                samples = Variable(samples, requires_grad=True)
                labels = Variable(labels, requires_grad=False)

                def closure(backward=True):
                    if torch.is_grad_enabled() and backward:
                        optimizer.zero_grad()
                    model_outputs = model(samples)
                    cri_loss = criterion(model_outputs, labels)
                    del model_outputs
                    if cri_loss.requires_grad and backward:
                        cri_loss.backward(retain_graph=True, create_graph=True)
                    return cri_loss

                # iter_timer.start()
                tr_loss = optimizer.step(closure=closure)

                # step_times.append(iter_timer.stop(tag='Iteration execution time', verbose=False))
                batch_loss = tr_loss.item()
                running_loss += batch_loss * train_loader.batch_size
                running_samples += train_loader.batch_size


                n_iter += 1
                if n_iter % 100 == 0:
                    with torch.no_grad():
                        correct = 0
                        total = 0
                        for tst_samples, tst_labels in test_loader:
                            tst_samples = Variable(tst_samples)
                            tst_outputs = model(tst_samples)
                            _, predicted = torch.max(tst_outputs.data, 1)
                            total += tst_labels.size(0)
                            correct += (predicted == tst_labels).sum()
                        accuracy = 100 * correct / total
                        print("n_iteration: {}. Tr. Loss: {}. Tst. Accuracy: {}.".format(n_iter,
                                                                                         running_loss / running_samples,
                                                                                         accuracy))

                progressbar.update(labels.size(0))

        tr_losses.append(running_loss / len(train_loader.sampler))

total_timer.stop(tag='Run execution time', verbose=True)
# tr_losses = preprocess(tr_losses)
import numpy as np
# plt.semilogy(np.cumsum(step_times), tr_losses)
plt.plot(tr_losses, 'o-')
# plt.yscale('log')
plt.grid()

plt.show()
</code></pre>
<p>if__name__ == ‘<strong>main</strong>’:<br>
batch_size = 32<br>
epochs = 5<br>
depth = 18</p>
<pre><code>task_name = 'cifar10'

if 'cifar10' in task_name:
    transform_train = transforms.Compose([
        transforms.RandomCrop(32, padding=4),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
    ])

    train_dataset = pydsets.CIFAR10(root='../../../Datasets', train=True, download=False, transform=transform_train)
    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=2)

    transform_test = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
    ])
    test_dataset = pydsets.CIFAR10(root='../../../Datasets', train=False, download=False, transform=transform_test)
    test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=2)

else:
    raise Exception('Select an available optimizer')

num_features = NUM_FEATURES[task_name]

main()
</code></pre>
</blockquote>

username: ptrblck
cooked: <p>Could you why <code>retain_graph=True, create_graph=True</code> are used?<br>
I don’t see a corresponding code snippet clearing the computation graph, so would expect to see an increased memory usage.</p>

username: ITs
cooked: <p>Thanks, <a class="mention" href="/u/ptrblck">@ptrblck</a> that was it. Isn’t <code>optimizer.zero_grad()</code> responsible for cleaning up? Could you please provide a reference to get some idea of what you mean?</p>

username: ptrblck
cooked: <p>No, <code>optimizer.zero_grad()</code> will set the <code>.grad</code> attributes of all passed parameters to zeros (or to <code>None</code> if <code>set_to_none=True</code> is used). The computation graph (and thus the intermediate tensors) will be freed during the <code>backward()</code> operation in the default setup (i.e. if <code>retain_graph=False</code> is used).</p>