[PyTorch] 训练的时候很好测试的时候爆显存
2018-12-15 本文已影响11人
VanJordan
- 先设置不使用梯度,然后将测试时候的batchsize设置成训练时候的二分之一或者三分之一就不会爆了。
- 可能原因是测试的时候真的需要更大的显存。
def early_test(self):
with torch.no_grad():
self.model.eval()
for path in tqdm(self.save_path,desc='Test Loop:'):
self.eval(path)
self.test()