PyTorch из памяти GPU сразу после инициализации

Я построил свою модель в PyTorch. Код хорошо работает на процессоре. Однако, похоже, что он исчерпывает память GPU сразу после инициализации сети и переключения на cuda.

Код ошибки:

model = init_from_scratch(args, train_exs, dev_exs)
model.init_optimizer()

nParams= sum([np.prod(list(p.size())) for p in model.network.parameters()])
print('* total number of parameters:',nParams)

device = torch.device("cuda:"+str(args.gpu) if args.cuda else "cpu")
model.set_device(device) #call last!!

# DATA ITERATORS  (not run to here)
train_dataset = data.ReaderDataset()
train_sampler = torch.utils.data.sampler.RandomSampler()
train_loader = torch.utils.data.DataLoader()
...


# TRAIN (not run to here)
for epoch in range(start_epoch, args.num_epochs):
....

def set_device(self,device):
    print("device",device)
    self.use_cuda = False if str(device) == "cpu" else True
    self.network = self.network.to(device)
  • общее количество параметров: 30708906

  • Версия факела - 0.4.0

  • TITAN Xp с 12G GPU памятью и NVIDIA-SMI показывает его бесплатно

Итак, что не так с моим кодом?

Вывод ошибки:

Traceback (most recent call last):
File "script/train2.0.py", line 683, in <module>
  main(args)
File "script/train2.0.py", line 542, in main
  model.set_device(device)
File "./model.py", line 588, in set_device
  self.network = self.network.to(device)
File "/home/username/.local/lib/python3.5/site-packages/torch/nn/modules/module.py", line 393, in to
  return self._apply(lambda t: t.to(device))
File "/home/username/.local/lib/python3.5/site-packages/torch/nn/modules/module.py", line 176, in _apply
  module._apply(fn)
File "/home/username/.local/lib/python3.5/site-packages/torch/nn/modules/module.py", line 182, in _apply
  param.data = fn(param.data)
File "/home/username/.local/lib/python3.5/site-packages/torch/nn/modules/module.py", line 393, in <lambda>
  return self._apply(lambda t: t.to(device))
File "/home/username/.local/lib/python3.5/site-packages/torch/cuda/__init__.py", line 161, in _lazy_init
  torch._C._cuda_init()

RuntimeError: ошибка времени выполнения cuda (2): недостаточно памяти в /pytorch/aten/src/THC/THCTensorRandom.cu:25

0 ответов

Другие вопросы по тегам