Pytorch-モデルをgpuまたはcpuにロード
840 ワード
多くの場合、gpuでモデルを訓練しますが、inferenceではgpuを使用したくないです.あるいは別のgpuで使いたいのですが、どうすればいいのでしょうか.
loadの時にdeviceを選択する必要があります.
モデルのパラメータ(model.state_dict()をファイルmodelに保存します.pth中.
1、cpu->cpuまたはgpu->gpu
この状況は最も簡単です.
2、cpu->gpu1
3、gpu 0 -> gpu 1
4、gpu -> cpu
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)