Pytorch——將模型load到gpu或cpu上

很多時候我們在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)

 

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章