Как мне избежать увеличения зарезервированной памяти CUDA? - PullRequest
Новая
       81

Как мне избежать увеличения зарезервированной памяти CUDA?

0 голосов
/ 02 августа 2020

Кто-нибудь может мне помочь, пожалуйста?)

Я пытаюсь обучить свою ConvNetwork в Google Colab с помощью PyTorch на torchvision.datasets.STL10 (10'000 примеров в поезде, 1'500 в наборе val, 1 '500 in test), но у меня снова и снова возникает одна и та же проблема - сеть начала обучаться, через несколько эпох процесс сломался, и появилась ошибка.

Когда я запускаю код ниже, что-то идет не так и объем памяти, который я использую, увеличивается, и через некоторое время у меня совсем не остается свободной памяти.

Понятия не имею, что не так. Я попытался понять это, но не смог найти решение своей проблемы даже в stackOverflow.

Я пытался использовать torch.cuda.empty_cache() в надежде, что он очистит память, которую использует CUDA, но ничего не изменилось.

Нейронная сеть: 

class ResNet(torch.nn.Module):
  def __init__(self):
    super(ResNet, self).__init__()
    
    self.conv1 = torch.nn.Conv2d(
        in_channels = 3, out_channels=30, kernel_size=5, padding=1) # 96x96x3 -> 94x94x30
    self.act1  = torch.nn.ReLU()
    self.pool1 = torch.nn.MaxPool2d(kernel_size=2, stride=2) # 94x94x30 -> 47x47x30

    self.conv2 = torch.nn.Conv2d(
        in_channels=30, out_channels=100, kernel_size=5, padding=1) # 47x47x30 -> 45x45x100
    self.act2  = torch.nn.ReLU()
    self.pool2 = torch.nn.MaxPool2d(kernel_size=3, stride=3) # 45x45x100 -> 15x15x100

    self.fc3  = torch.nn.Linear(in_features=15*15*100, out_features = 1000)
    self.act3 = torch.nn.ReLU()

    self.fc4  = torch.nn.Linear(in_features=1000, out_features=1000)
    self.act4 = torch.nn.ReLU()

    self.fc5  = torch.nn.Linear(in_features=1000, out_features=10)


  def forward(self, x):

    x = self.conv1(x)
    x = self.act1(x)
    x = self.pool1(x)

    x = self.conv2(x)
    x = self.act2(x)
    x = self.pool2(x)
    
    x = x.reshape(x.shape[0], 22500)

    x = self.fc3(x)
    x = self.act3(x)

    x = self.fc4(x)
    x = self.act4(x)

    x = self.fc5(x)
    return x

Следующий сценарий: 

NNet = ResNet()

device = torch.device('cuda:0' if torch.cuda.is_available else 'cpu')
loss = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(NNet.parameters(), lr = 0.001)
NNet = NNet.to(device)

random.seed(k)
np.random.seed(k)
torch.manual_seed(k)
torch.cuda.manual_seed(k)
torch.backends.cudnn.deterministic = True

X_train = X_train.to(device)
y_train = y_train.to(device)

X_dev = X_dev.to(device)
y_dev = y_dev.to(device)

Сценарий обучения: 

batch_size = 1000

for epoch in range(5):
  
  optimizer.zero_grad()

  order = np.random.permutation(X_train.shape[0])
  for start_batch in range(0, X_train.shape[0], batch_size):

    optimizer.zero_grad()
    batch_index = order[start_batch:start_batch+batch_size]
    X_batch = X_train[batch_index].to(device)
    y_batch = y_train[batch_index].to(device)

    y_pred = NNet.forward(X_batch)

    loss_train = loss(y_pred, y_batch)
    loss_train.backward()
    optimizer.step()

  accuracy_train = (y_pred.argmax(1) == y_batch).data.float().mean().cpu() # Train Accuracy
  y_pred = NNet.forward(X_dev)
  loss_dev = loss(y_pred, y_dev).data.cpu()
  accuracy_dev = (y_pred.argmax(dim=1) == y_dev).data.float().mean().cpu() # Dev Accuracy

  print('epoch #%i' % epoch, end='\t')
  print('loss dev: %.4f' % loss_dev, end='\t')
  print('loss train: %.4f' % loss_train, end='\t')
  print('accuracy dev: %.4f' % accuracy_dev, end='\t')
  print('accuracy train: %.4f' % accuracy_train)

RuntimeError: CUDA не хватает памяти. Пытался выделить 774,00 МиБ (GPU 0; 11,17 ГиБ общая емкость; 10,48 ГиБ уже выделено; 275,81 МиБ бесплатно; 10,54 ГиБ зарезервировано всего PyTorch

Всем спасибо!

...