很多時候我們在gpu上訓練一個模型,但是在inference的時候不想使用gpu。或者想在別的gpu上使用,那麼怎麼辦呢?
需要在load的時候就選擇device。
保存了模型的參數(model.state_dict())到文件model.pth中。
1、cpu->cpu 或gpu->gpu
這種情況是最簡單的:
checkpoint = torch.load('model.pth')
model.load_state_dict(checkpoint)
2、cpu->gpu1
checkpoint=torch.load('model.pth', map_location=lambda storage, loc: storage.cuda(1))
model.load_state_dict(checkpoint)
3、gpu 0 -> gpu 1
checkpoint=torch.load('model.pth', map_location={'cuda:0':'cuda:1'})
model.load_state_dict(checkpoint)
4、gpu -> cpu
checkpoint=torch.load('model.pth', map_location=lambda storage, loc: storage)
model.load_state_dict(checkpoint)