MobileNet学習中にTrain関数は動作するが、Validate関数内でCUDA out of memory とエラーが出る

Question

前提・実現したいこと

前回の質問から引き続き、MobileNet V1を実装中。
学習データはImageNetのスモールデータセットであるImageNetteを用いる。
このデータセットは

!wget https://s3.amazonaws.com/fast-ai-imageclas/imagenette-320.tgz
!tar xzf imagenette-320.tgz

コマンドで取得しColab上から対象のGoogleドライブ内のフォルダにアクセスし読み込んでいる。

ベースとなる学習コードはGitHub上の

https://github.com/Z0m6ie/CIFAR-10_PyTorch

を参考にしている。

このデータセットを用いて学習を行いたい。

発生している問題・エラーメッセージ

Train関数(後述)での訓練が終了しValidate関数(後述)を実行時に以下のエラーが発生する。

/usr/local/lib/python3.6/dist-packages/ipykernel_launcher.py:15: UserWarning: volatile was removed and now has no effect. Use `with torch.no_grad():` instead.
  from ipykernel import kernelapp as app
/usr/local/lib/python3.6/dist-packages/torch/nn/_reduction.py:43: UserWarning: size_average and reduce args will be deprecated, please use reduction='sum' instead.
  warnings.warn(warning.format(ret))
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-83-f5e663fb5c57> in <module>()
      9 for epoch in range(80):
     10     #train(epoch)
---> 11     loss, accuracy = validate(epoch)
     12     best_loss, best_acc = save_best(loss, accuracy, best_loss, best_acc)
     13 

7 frames
/usr/local/lib/python3.6/dist-packages/torch/nn/modules/conv.py in _conv_forward(self, input, weight)
    344                             _pair(0), self.dilation, self.groups)
    345         return F.conv2d(input, weight, self.bias, self.stride,
--> 346                         self.padding, self.dilation, self.groups)
    347 
    348     def forward(self, input):

RuntimeError: CUDA out of memory. Tried to allocate 196.00 MiB (GPU 0; 7.43 GiB total capacity; 6.80 GiB already allocated; 100.94 MiB free; 6.82 GiB reserved in total by PyTorch)

1epoch回すだけで1時間学習にかかるのでこのエラーは正直つらいです。

該当のソースコード

前回同様、学習データのロード,Train関数,Validate関数,それら関数を実行するコードを記す。

import argparse
import torch
import numpy as np
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.optim.lr_scheduler import StepLR
from torchvision import datasets, transforms
from torch.autograd import Variable
from torch.utils.data.sampler import SubsetRandomSampler
from sklearn.metrics import accuracy_score
from torchsummary import summary

use_cuda = torch.cuda.is_available()
dtype = torch.cuda.FloatTensor if use_cuda else torch.FloatTensor

normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                 std=[0.229, 0.224, 0.225])

valid_size=0.1

traindir = 'imagenette-320/train'
valdir = 'imagenette-320/val'

train_transform = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    normalize
])

valid_transform = transforms.Compose([
         transforms.Resize(256),
         transforms.CenterCrop(224),
         transforms.ToTensor(),
         normalize
])

train_dataset = datasets.ImageFolder(traindir,transform=train_transform)
valid_dataset = datasets.ImageFolder(valdir,transform=valid_transform)

train_sampler = None
valid_sampler = None


train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=128,sampler=train_sampler,shuffle=True)

valid_loader = torch.utils.data.DataLoader(valid_dataset, batch_size=128,sampler=valid_sampler,shuffle=True)

def train(epoch):
    model.train()
    writer = SummaryWriter()
    for batch_idx, (data, target) in enumerate(train_loader):
        if use_cuda:
            data, target = data.cuda(), target.cuda()
        data, target = Variable(data), Variable(target)
        optimizer.zero_grad()
        output = model(data)
        correct = 0
        pred = output.data.max(1, keepdim=True)[1] # get the index of the max log-probability
        correct += pred.eq(target.data.view_as(pred)).sum()
        loss = criterion(output, target)
        loss.backward()
        accuracy = 100 * (int(correct)/ int(len(output)))
        #accuracy = 100. * correct / len(output)
        optimizer.step()
        if batch_idx % 1 == 0:
             print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}, Accuracy: {:.2f}'.format(epoch, batch_idx * len(data), len(train_loader.dataset),100. * batch_idx / len(train_loader), loss.data.item(), accuracy))
             writer.add_scalar('Loss/Loss', loss.data.item(), epoch)
             writer.add_scalar('Accuracy/Accuracy', accuracy, epoch)
    scheduler.step()

def validate(epoch):
    valid_loss=0.0
    accuracy=0.0
    model.eval()
    
    writer = SummaryWriter()
    
    valid_loss = 0
    correct = 0

    for data, target in valid_loader:
        if use_cuda:
            data, target = data.cuda(), target.cuda()
        
        data, target = Variable(data, volatile=True), Variable(target)
        output = model(data)
        valid_loss += F.cross_entropy(output, target, size_average=False).data.item() # sum up batch loss
        pred = output.data.max(1, keepdim=True)[1] # get the index of the max log-probability
        correct += pred.eq(target.data.view_as(pred)).sum()
        
    
    valid_loss /= len(valid_loader.dataset)
    accuracy = 100. * correct / len(valid_loader.dataset)
    print('\nValidation set: Average loss: {:.4f}, Accuracy: {}/{} ({:.2f}%)\n'.format(valid_loss, correct, len(valid_loader.dataset),100. * correct / len(valid_loader.dataset)))
    writer.add_scalar('Loss/Validation_Loss', valid_loss, epoch)
    writer.add_scalar('Accuracy/Validation_Accuracy', accuracy, epoch)
    
    return valid_loss, accuracy

!pip install tensorboardX
from tensorboardX import SummaryWriter
best_loss = None
best_acc = None

for epoch in range(80):
    #train(epoch)
    loss, accuracy = validate(epoch)
    best_loss, best_acc = save_best(loss, accuracy, best_loss, best_acc)
    
writer.export_scalars_to_json("./all_scalars.json")

writer.close()

試したこと

エラーメッセージの前に警告が出ていた。その警告にを知らべると自分と同じようなエラーが出ている人が他にもいたので参考に,**with torch.no_grad():**をValidate関数内に以下のように付け足した。

def validate(epoch):
    valid_loss=0.0
    accuracy=0.0
    model.eval()
    
    writer = SummaryWriter()
    
    valid_loss = 0
    correct = 0
    with torch.no_grad():
      for data, target in valid_loader:
          if use_cuda:
              data, target = data.cuda(), target.cuda()
          
          data, target = Variable(data, volatile=True), Variable(target)
          output = model(data)
          valid_loss += F.cross_entropy(output, target, size_average=False).data.item() # sum up batch loss
          pred = output.data.max(1, keepdim=True)[1] # get the index of the max log-probability
          correct += pred.eq(target.data.view_as(pred)).sum()
          
      
      valid_loss /= len(valid_loader.dataset)
      accuracy = 100. * correct / len(valid_loader.dataset)
      print('\nValidation set: Average loss: {:.4f}, Accuracy: {}/{} ({:.2f}%)\n'.format(valid_loss, correct, len(valid_loader.dataset),100. * correct / len(valid_loader.dataset)))
      writer.add_scalar('Loss/Validation_Loss', valid_loss, epoch)
      writer.add_scalar('Accuracy/Validation_Accuracy', accuracy, epoch)
      
    return valid_loss, accuracy

警告文は消えたが、同様のエラーが出力された。また、Variable(target)に引数volatile=Trueを入れてみても同様の結果であった。
どなたか原因と対策をご教授頂ければ幸いです。

補足情報（FW/ツールのバージョンなど）

ここにより詳細な情報を記載してください。

Answer 1

CUDA out of memory は GPU メモリが足りないというエラーです。なので、今のままでは実行できないので、ハードを変えられないのであれば、使用するメモリ量を減らす必要があります。
必要なメモリ量は、画像の入力サイズ、バッチサイズ、モデルの種類などの要素で決まります。
Colab で GPU メモリがどのくらい使えるのかわかりませんが、バッチサイズ batch_size=128 というのはかなり大きいので、学習用だけで GPU メモリが専有されて、バリデーション時に追加で GPU メモリを確保しようとした際にメモリが足りないとなっているのではないでしょうか。
batch_size を64, 32, 16のようにメモリが足りないとエラーにならない程度まで減らしていけばよいと思います。

with torch.no_grad():をValidate関数内に以下のように付け足した。

推論時はバックプロパゲーションをしないので、中間層の出力を記録しておく必要がないので、これは有効です。

1epoch回すだけで1時間学習にかかるのでこのエラーは正直つらいです。

学習のほうを1イテレーションだけ回したあとにすぐ break してバリデーションすれば、今の設定でメモリ不足かどうかはすぐに確認できると思います。

Comments

[Otom_jp]

ご回答ありがとうございます。
訓練時は同様のバッチサイズで問題なかったのに、検証時にメモリが足りないということはあり得るのでしょうか。

[tiitoi]

no_grad() を指定していれば、推論時はあまりメモリを使用しないので、推論するタイミングで足りなくなるということは普通あまりないのですが、本当にメモリ使用量が逼迫しているのであれば、足りないということはありえます。

とりあえず、バッチサイズを小さくして、学習は1エポック回さなくても1回だけループしてすぐ break し、バリデーションしてエラーにならないかどうか確認してみてください
バッチサイズをかなり小さくしても以前としてエラーになるのであれば、別の原因の可能性もありますが。。

[tiitoi]

それか Google Colab は割り当てられる GPU の性能が毎回異なると聞いたことがあるので、性能がいい GPU が引けるまで、頑張ってみるといいかもしれません。

https://qiita.com/koshian2/items/d33edc963ed6cfcad77e

[Otom_jp]

ご回答ありがとうございます。
バッチサイズを64に変更したところ正常にValidate関数が動作し、訓練時の1epochにかかる時間も大幅に改善されました。ありがとうございます。助かりました。