问题描述
在使用Pytorch跑模型时,虽然用了GPU,但随着训练的进行,后台占用内存(不是显存)越来越多,直到内存不足,程序崩溃:
OSError: [Errno 12] Cannot allocate memory
解决办法
造成这个问题的原因是我在使用循环计算一个epoch内所有batch的loss之和时,直接累加了tensor,导致梯度信息一直保持,内存占用不断增加:
for epoch in range(1, n_epoch+1):
loss_epoch = 0
model.train()
for batch in loader:
......
loss = criterion(......)
loss.backward()
optimizer.step()
loss_epoch += loss # All loss in a batch
修改方法是,loss先转成浮点数,再累加:
for epoch in range(1, n_epoch+1):
loss_epoch = 0
model.train()
for batch in loader:
......
loss = criterion(......)
loss.backward()
optimizer.step()
loss_epoch += loss.item() # All loss in a batch