[Pytorch]cuda out of memory when increase loop iteration

今天發現了這陣子一直沒有解決的一個bug，蠢到都想挖一個洞把自己埋起來，還是記錄一下以免日後又犯蠢。

不太清楚要怎麼用英文下這個標題，主要的問題敘述是：在使用pytorch訓練模型的時候，為了紀錄模型對於testing set的效果，每隔幾個epoch我就會將acc / loss append到一個list中，然後我的模型常常會發生一個很神奇的狀況：

剛開始跑的時候GPU mem還塞得下，但隨著loop次數增加，mem居然也相對的增加了，導致幾個loop後就會發生”cuda out of memory”的問題。

一開始我還以為是模型太大，可是後來想想，不對啊，如果前幾個loop是可以跑完的，那就代表gpu memoey應該是可以塞下整個模型的，後來去檢查我的code和上網搜尋相關資源後才發現我犯了一個很蠢的事情，先給大家看一下我的code大概是怎麼寫的：

for e in range(epoch):
    # training
    model.train()
    for batch_idx, (batch_X, batch_y) in enumerate(train_loader):
        batch_X = batch_X.to(device)
        batch_y = batch_y.to(device)
        output, _ = model(batch_X)
        loss = criterion(output, batch_y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
 
    # testing
    if e % 10 == 0:
        model.eval()
        X_test = X_test.to(device)
        y_test = y_test.to(device)
        output, _ = model(X_test)
        # get the class of prediction
        prediction = torch.max(output, dim=1)[1]
        loss = criterion(output, y_test)
        acc = torch.sum(prediction == y_test).cpu().numpy() / X_test.shape[0]
        acc_history.append(acc)
        loss_history.append(loss)
        print('Epoch: {epoch} | loss: {loss} | acc:{acc}'.format(epoch=e, loss=loss, acc=acc))
        torch.save(model.state_dict(), model_weight_filepath)

問題主要出在22~24行，我會把每次model對testing set的acc和loss記錄下來，可是我沒有先將這些tensor variable進行detach()，所以這些變數其實還包含了computation graph的資訊，造成每次append到list的時候都會重複佔用memory。

解決方法就是將不需要backward的variables都detatch()，也就是將22~24行的code改成：

1
2
3

acc = torch.sum(prediction == y_test).detach().cpu().numpy() / X_test.shape[0]
acc_history.append(acc)
loss_history.append(loss.detach())

如此一來，迴圈越跑越多memory也不會越來越大了！網路上也有人遇到一樣的問題：可以看看原文I run out of memory after a certain amount of batches when training a resnet18