1. 问题
训练模型时报错:
RuntimeError: CUDA out of memory. Tried to allocate 128.00 MiB (GPU 0; 2.00 GiB total capacity; 1.49 GiB already allocated; 57.03 MiB free; 6.95 MiB cached)
2. 分析
这种问题,是GPU内存不够引起的
3. 解决
- 方法一:通过修改代码减少GPU内存使用量
(1)不计算梯度
报错的训练代码为:
def train_model(model, criterion, optimizer, dataload, num_epochs=20):
for epoch in range(num_epochs):
print('Epoch {}/{}'.format(epoch, num_epochs - 1))
print('-' * 10)
dt_size = len(dataload.dataset)
epoch_loss = 0
step = 0
for x, y in dataload:
step += 1
inputs = x.to(device)
labels = y.to(device)
# zero the parameter gradients
optimizer.zero_grad()
# forward
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()#梯度下降,计算出梯度
optimizer.step()
epoch_loss += loss.item()
print("%d/%d,train_loss:%0.3f" % (step, (dt_size - 1) // dataload.batch_size + 1, loss.item()))
print("epoch %d loss:%0.3f" % (epoch, epoch_loss/step))
torch.save(model.state_dict(), 'weights_%d.pth' % epoch)# 返回模型的所有内容
return model
在 outputs = model(inputs)
那里加上一行代码:with torch.no_grad():
...
optimizer.zero_grad()
# forward
with torch.no_grad():#添加这行代码
outputs = model(inputs)#再缩进这行
loss = criterion(outputs, labels)
...
说明:对于tensor的计算操作,默认是要计算梯度和进行反向传播的,而torch.no_grad()用来禁止梯度的计算和反向传播,可以减少GPU的内存使用量。
在with torch.no_grad():
作用域范围内的操作不会计算梯度,也不会进行反向传播。
如果添加上述代码后还不行,那是因为我们前面虽然禁止了计算梯度,但是却没有计算梯度的Variable,因为这里的inputs本身默认就是不求梯度的。所以还会报错:
RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn
这时需要构建Variable,并传入参数:requires_grad=True
,这个参数表示对这个变量求梯度,默认是False。
...
inputs = x.to(device)
labels = y.to(device)
#添加下面两行代码
inputs = Variable(inputs,requires_grad=True)
labels = Variable(labels,requires_grad=True)
...
记得还需要在代码最前面添加 from torch.autograd import Variable
才能导入Variable
(2)调整Batch Size
还有一种方法就是调整
Batch Size
(即一次训练所抓取的数据样本数量)的大小,Batch size大小的选择也至关重要。为了在内存效率和内存容量之间寻求最佳平衡,batch size应该精心设置,从而最优化网络模型的性能及速度。
在卷积神经网络中,大的batch size通常可使网络更快收敛,但由于内存资源的限制,batch size过大可能会导致内存不够用,如果GPU内存不够用出现以上情况,可以将Batch Size适当的设置小一点。
- 方法二:更换更高性能更高显存的显卡
这样所有问题都解决了。