Я написал код для обучения модели глубокого обучения, где я удаляю туды cuda после каждой партии и выполняю torch.cuda.empty_cache()
после этого.Я уверен, что размер пакета недостаточно велик, чтобы вызвать эту ошибку.Что может быть причиной этого?
for epoch in range(1+last_epoch, self.num_epochs+1):
for phase in ['train', 'val']:
loss_arr = []
if phase == 'train':
model.train()
scheduler.step()
was_training = True
else:
model.eval()
was_training = False
for i_batch, sample_batched in enumerate(dataloaders[phase]):
X = sample_batched[0]
y = sample_batched[1].type(torch.LongTensor)
w = sample_batched[2]
if model.is_cuda:
X, y, w = X.cuda(non_blocking=True), y.cuda(non_blocking=True), w.cuda(non_blocking=True)
output = model(X)
loss = self.loss_func(output, y, w)
if phase == 'train':
curr_iteration+=1
optim.zero_grad()
loss.backward()
optim.step()
if (curr_iteration % log_nth == 0):
self.logWriter.loss_per_iter(loss.item(), curr_iteration)
loss_arr.append(loss.item())
with torch.no_grad():
self.logWriter.update_cm_per_iter(output, y, self.labels, phase)
del X, y, w, output, loss
torch.cuda.empty_cache()
self.logWriter.loss_per_epoch(loss_arr, phase, epoch)
epoch_output, epoch_labels = model.predict(dataloaders[phase].dataset.X), dataloaders[phase].dataset.y
self.logWriter.dice_score_per_epoch(epoch_output, epoch_labels, phase, epoch)
index = np.random.choice(len(dataloaders[phase].dataset), 3, replace=False)
self.logWriter.image_per_epoch(epoch_output[index], epoch_labels[index], phase, epoch)
self.logWriter.cm_per_epoch(self.labels, phase, epoch, i_batch)
del epoch_output, epoch_labels
print("==== Epoch ["+str(epoch)+" / "+str(self.num_epochs)+"] done ====")
model.save('models/' + self.exp_dir_name + '/quicknat_epoch' + str(epoch) + '.model')
И функция perdict внутри модели выглядит следующим образом
def predict(self, X, enable_dropout = False):
"""
Predicts the outout after the model is trained.
Inputs:
- X: Volume to be predicted
"""
self.eval()
if type(X) is np.ndarray:
X = torch.tensor(X, requires_grad = False).cuda(non_blocking=True)
elif type(X) is torch.Tensor and not X.is_cuda:
X = X.cuda(non_blocking=True)
if enable_dropout:
self.enable_test_dropout()
with torch.no_grad():
out = self.forward(X)
max_val, idx = torch.max(out,1)
idx = idx.data.cpu().numpy()
prediction = np.squeeze(idx)
del X, out, idx, max_val
return prediction