问题描述:如果训练代码中不加入验证阶段。程序能够正常执行,不会出现下面cuda超出的现象。
但是,如果用epoch=1阶段结束后的模型去执行验证阶段,会出现验证集的第一个banch不会报错,但是执行第二个batch时,会提示cuda超了。原因可能是由于梯度在验证阶段不回传,造成梯度的累计。
解决办法:
把生成代码加入with torch.no_grad():内,来取消验证阶段的loss
with torch.no_grad(): # 由于梯度在验证阶段不回传,会造成梯度的累计,所以需要加上这句
gen_imgs = G(torch.autograd.Variable(label.cuda()))
验证代码:
# 验证:从样本集中选一些数据做验证
epoch_path = os.path.join(save_dir, "val_epoch%s" %(epoch))
visualize.mkdirs(epoch_path)
for n, items in enumerate(valDataset):
# input image data
label, _, name = items
with torch.no_grad(): # 由于梯度在验证阶段不回传,会造成梯度的累计,所以需要加上这句
gen_imgs = G(torch.autograd.Variable(label.cuda()))
visualize.saveImg(gen_imgs, epoch_path, 'generated', name[0])
message = "val: iter %d/%d" % (n + 1, valDatset_size)
print(message)