如题这是经常遇见的显存不够问题,一般分为train阶段与valid阶段,我曾经在valid阶段时因为没有进行手动的梯度清零而导致显存不够,加上with t.no_grad:问题迎刃而解。
with t.no_grad():
for ii, data in enumerate(dataloader):
input, label = data
val_input1 = Variable(input, volatile=True)
val_input = val_input1.float()
val_label = Variable(label.type(t.LongTensor), volatile=True)
if opt.use_gpu:
val_input = val_input.cuda()
val_label = val_label.cuda()
score = model(val_input)
loss = criterion(score,val_label) / (160 * 160)
if ii == 0:
loss_mean = loss.item() # 通过.item()返回python number
loss_mean = 0.1 * loss + 0.9 * loss_mean
之后当我对原有程序进行更改以及重新定义dataloder后,在train的第一个epoch中就出现了显存不够问题具体报错如下:
File "d:\anaconda3\lib\site-packages\fire\core.py", line 138, in Fire
component_trace = _Fire(component, args, parsed_flag_args, context, name)
File "d:\anaconda3\lib\site-packages\fire\core.py", line 468, in _Fire
target=component.__name__)
File "d:\anaconda3\lib\site-packages\fire\core.py", line 672, in _CallAndUpdateTrace
component = fn(*varargs, **kwargs)
File "main.py", line 103, in train
optimizer.step()
File "d:\anaconda3\lib\site-packages\torch\autograd\grad_mode.py", line 15, in decorate_context
return func(*args, **kwargs)
File "d:\anaconda3\lib\site-packages\torch\optim\adam.py", line 107, in step
denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(group['eps'])
可以看出实在第一次梯度更新的时候显存直接爆炸。而我训练阶段的代码如下:
for epoch in range(opt.max_epoch):
loss_mean = 0.0
loss_val = 0.0
loss_meter.reset()
for ii,(data,label) in tqdm(enumerate(train_dataloader),total=len(train_data)):
# train model
input = Variable(data)
input = input.float()
target = Variable(label)
if opt.use_gpu:
input = input.cuda()
target = target.cuda()
optimizer.zero_grad()
score = model(input)
loss = criterion(score,target) / (400 * 190)
loss.backward()
optimizer.step()
最后借助论坛大佬的帮助成功解决问题:
# train model
input = Variable(data)
input = input.float()
target = Variable(label)
if opt.use_gpu:
input = input.cuda()
target = target.cuda()
optimizer.zero_grad()
with torch.no_grad() :
score = model(input)
loss = criterion(score,target) / (400 * 190)
loss = Variable(loss, requires_grad = True)
loss.backward()
optimizer.step()
具体思路就是将score与loss不进行梯度更新、然后将loss赋给固定变量从而防止重复赋值占用显存。
不得不说,现有框架用着实在太蠢了,同一个程序在第一次跑的时候train不会出现这种问题,第二次就出现了。
torch.Size([2, 400, 190])
----------------------------------------------------------------
Layer (type) Output Shape Param #
================================================================
Conv2d-1 [-1, 64, 384, 160] 1,216
MaxPool2d-2 [-1, 64, 192, 80] 0
Conv2d-3 [-1, 128, 192, 80] 73,856
MaxPool2d-4 [-1, 128, 96, 40] 0
Conv2d-5 [-1, 256, 96, 40] 295,168
Conv2d-6 [-1, 256, 96, 40] 590,080
MaxPool2d-7 [-1, 256, 48, 20] 0
Conv2d-8 [-1, 512, 48, 20] 1,180,160
Conv2d-9 [-1, 512, 48, 20] 2,359,808
MaxPool2d-10 [-1, 512, 24, 10] 0
Conv2d-11 [-1, 512, 24, 10] 2,359,808
Conv2d-12 [-1, 512, 24, 10] 2,359,808
MaxPool2d-13 [-1, 512, 12, 5] 0
Dropout2d-14 [-1, 512, 12, 5] 0
Linear-15 [-1, 19000] 583,699,000
Linear-16 [-1, 19000] 361,019,000
更新:在train的时候是不能使用with torch.no_grad.会导致模型参数无法更新。我的模型出现问题是由于更改了输入输出导致计算量巨大可以见到在两层全连接层总计算参数达到了9亿,这还是在batchsize为1的情况下预计的最小计算参数。因此这也就导致了纵使我改低BS,32G显卡下跑仍然会造成显存溢出的情况。因此后续我将会将网络修改为全卷积网络来测试出现能否正常运行。
2020.11.17更新:利用全卷积网络并不能降低参数量。因而想到另外几个方法:
- 能不能将fc1的output units换为4096,fc2的用19000
- 若第一个方案不行的话,换成全卷积网络,其中替代FC层的卷积层换成深度可分离卷积。
- 另一个方案为对图像进行降采样,将图像降为200*90.这样就能大大降低参数量了。不过同样这个方案会降低分辨率。
- 另有一方法全局平均池化,在论文中说能替代FC层,且替代后效果更好。论文:Network In Network。
2020.11.19更新:运用GAP层成功解决问题。大大缩减了参量。详情使用请见我另一篇博客GAP层
2020.11.26更新:由于我使用GAP层窗口过大导致泛化误差降不下来,因此我将输出进行降采样使用回了FC层。泛化误差得以正常。因此使用GAP层的时候注意窗口不能太大,同样我建议使用全局最大池化层,不会出现GAP一样的进行了均值滤波的效果。
论坛问题地址:torch forum