torch.load CUDA ERROR: out of memory

在pytorch中load模型，出現CUDA ERROR:out of memory,基本有下面三個原因：

CUDA顯存不足

解決辦法：換成顯存更大的顯卡

顯卡被占用

當存在多個顯卡時，pytorch默認使用0號顯卡，當0號顯卡被占用時，即使1號顯卡空閑仍然會報錯
解決辦法：讓其他顯卡不可見，此時0號顯卡為空余的1號顯卡

import os

os.environ['CUDA_VISIBLE_DEVICE']='1'

模型保存GPU不同

使用cuda：0訓練保存的模型，如果直接使用cuda：1去load，也會報錯
解決辦法：model.load_state_dict(torch.load('trained.pth', map_location={'cuda:0': 'cuda:1'}))

參考