很多时候我们在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)