该文介绍了在PyTorch中加载模型的步骤，包括实例化模型如UNet，使用`load_state_dict`加载权重，以及在CPU或GPU上加载模型的注意事项。如果在CPU上加载，需指定`map_location`参数。
记录一下pytorch加载模型的一些细节
正确过程
import torch
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = UNet()
model.load_state_dict(torch.load('last.pth'))
model.to(device=device)
一些细节
1.注意模型的实例化UNet()
2.若在cpu加载模型，一般需要在torch.load中加上map_location=torch.device('cpu')，即
model.load_state_dict(torch.load('last.pth', map_location=torch.device('cpu')))
扫一扫
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
