Блог

Я пытаюсь построить модель автоэнкодера, где ввод/вывод-это изображения RGB размером 256 x 256. Я пытался обучить модель на 1 графическом процессоре с 12 ГБ памяти, но я всегда ловил CUDA OOM (я пробовал разные размеры пакетов, и даже размер пакета 1 не удается). Поэтому я прочитал о параллелизме моделей в Pytorch и попробовал это:

 class Autoencoder(nn.Module):  def __init__(self, input_output_size):  super(Autoencoder, self).__init__()  self.encoder = nn.Sequential(  nn.Linear(input_output_size, 1024),  nn.ReLU(True),  nn.Linear(1024, 200),  nn.ReLU(True)  ).cuda(0)    self.decoder = nn.Sequential(  nn.Linear(200, 1024),  nn.ReLU(True),  nn.Linear(1024, input_output_size),  nn.Sigmoid()).cuda(1)    print(self.encoder.get_device())  print(self.decoder.get_device())   def forward(self, x):  x = x.cuda(0)  x = self.encoder(x)  x = x.cuda(1)  x = self.decoder(x)  return x   

Поэтому я переместил свой кодер и декодер на разные графические процессоры. Но теперь я получаю это исключение:

 Expected tensor for 'out' to have the same device as tensor for argument #2 'mat1'; but device 0 does not equal 1 (while checking arguments for addmm)  

Это появляется, когда я делаю x = x.cuda(1) в прямом методе.

Более того, вот мой код «поезда», может быть, вы можете дать мне несколько советов по оптимизации? Являются ли изображения размером 3 x 256 x 256 слишком большими для обучения? (Я не могу их уменьшить). Заранее спасибо.

Обучение:

 input_output_size = 3 * 256 * 256 model = Autoencoder(input_output_size).to(device) optimizer = optim.Adam(model.parameters(), lr=1e-4) criterion = nn.MSELoss()  for epoch in range(100):  epoch_loss = 0  for batch_idx, (images, _) in enumerate(dataloader):  images = torch.flatten(images, start_dim=1).to(device)  output_images = model(images).to(device)  train_loss = criterion(output_images, images)      train_loss.backward()  optimizer.step()   if batch_idx % 5 == 0:  with torch.no_grad():  model.eval()  pred = model(test_set).to(device)  model.train()   test_loss = criterion(pred, test_set)   wandb.log({"MSE train": train_loss})  wandb.log({"MSE test": test_loss})  del pred, test_loss   if batch_idx % 200 == 0:  # here I send testing images from output to Wamp;B  with torch.no_grad():  model.eval()  pred = model(test_set).to(device)  model.train()  wandb.log({"PRED": [wandb.Image((pred[i].cpu().reshape((3, 256, 256)).permute(1, 2, 0) * 255).numpy().astype(np.uint8), caption=str(i)) for i in range(20)]})  del pred  gc.collect()  torch.cuda.empty_cache()  epoch_loss  = train_loss.item()  del output_images, train_loss  epoch_loss = epoch_loss / len(dataloader)  wandb.log({"Epoch MSE train": epoch_loss})  del epoch_loss  

Три проблемы, которые я вижу:

 model(test_set)  

Это происходит, когда вы отправляете весь свой набор тестов (предположительно огромный) в виде одной партии через свою модель.

Я не знаю, что wandb это такое, но еще одним вероятным источником роста памяти являются эти строки:

 wandb.log({"MSE train": train_loss}) wandb.log({"MSE test": test_loss})  

Вы , кажется, сохраняете train_loss и test_loss , но они содержат не только сами числа, но и вычислительные графики (живущие на графическом процессоре), необходимые для backprop. Прежде чем сохранять их, вы хотите преобразовать их в float или numpy .

Ваша модель содержит два 3*256*256 x 1024 блока веса. При использовании в Adam для них потребуется 3*256*256 x 1024 * 3 * 4 bytes = 2,25 ГБ VRAM каждый (возможно, больше, если это неэффективно реализовано). Это выглядит как плохая архитектура и по другим причинам.