pytorch中的torch.load怎么使用

692
2024/1/20 20:45:52
栏目: 编程语言
开发者测试专用服务器限时活动,0元免费领,库存有限,领完即止! 点击查看>>

在PyTorch中,torch.load()函数用于加载保存的模型或张量。其基本语法如下:

torch.load(filepath, map_location=None, pickle_module=<module 'pickle' from '...'>)
  • filepath是保存模型或张量的文件路径。
  • map_location是一个可选参数,用于指定设备将模型/张量加载到哪个位置。可以是一个字符串,表示设备名称(如’cpu’、'cuda:0’等),也可以是一个torch.device对象。默认值为None,表示加载到与保存时设备相同的位置。
  • pickle_module是一个可选参数,用于覆盖默认的pickle模块。默认值为Python内置的pickle模块。

以下是torch.load()函数的使用示例:

import torch

# 加载保存的模型
model = torch.load('model.pth')

# 加载保存的张量
tensor = torch.load('tensor.pt')

# 加载保存的模型,并将其加载到指定设备上
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
model = torch.load('model.pth', map_location=device)

# 加载保存的模型,使用自定义的pickle模块
import pickle5 as pickle
model = torch.load('model.pth', pickle_module=pickle)

注意,torch.load()函数只能加载在相同版本的PyTorch中保存的模型或张量。如果模型或张量是在不同版本的PyTorch中保存的,则需要使用其他方法进行转换或加载。

辰迅云「云服务器」,即开即用、新一代英特尔至强铂金CPU、三副本存储NVMe SSD云盘,价格低至29元/月。点击查看>>

推荐阅读: pytorch多线程任务咋分配