笔者近日利用pytorch框架对图像采用Resnet50模型提取特征的时候,利用到了GPU,发现报错如下:
RuntimeError: CUDA out of memory. Tried to allocate 858.00 MiB (GPU 3; 10.92 GiB total capacity; 10.10 GiB already allocated; 150.69 MiB free; 10.13 GiB reserved in total by PyTorch)
该程序运行所用GPU为1080TI板子,有大约11GB的显存,然而简单的Resnet50模型居然搞的显存不足。经过详细调试发现,需要使用 with torch.no_grad(): 语句,否则会因为存储过多中间数据而导致显存不足。
model=model.to(device)
with torch.no_grad():
img = Image.open(imgFile)
img = transform(img)
x = Variable(torch.unsqueeze(img, dim=0).float(), requires_grad=False)
x = x.to(device)
y = model(x).cpu()
尤其是在y=model(x).cpu()的时候,很容易出错。希望能对大家有所帮助,谢谢。