torch.load

torch.load 的基本用法

torch.load() 是 PyTorch 中用于加载由 torch.save() 保存的模型或张量的函数。它支持从文件路径或文件对象加载数据，并可以指定设备（CPU 或 GPU）。

import torch

# 从文件路径加载模型或张量
loaded_data = torch.load('path/to/file.pt')

常用参数说明

• f: 文件路径或文件对象（必需）。
• map_location: 指定加载的设备，如 'cpu' 或 'cuda:0'。
• pickle_module: 用于反序列化的 pickle 模块（默认 pickle）。
• weights_only: 是否仅加载权重（安全模式，默认 False）。

# 加载到指定设备
loaded_data = torch.load('model.pt', map_location=torch.device('cuda'))

加载模型权重

通常用于加载预训练模型的权重到现有模型结构：

model = MyModel()
model.load_state_dict(torch.load('weights.pt'))

加载整个模型

如果保存时包含模型结构和权重（通过 torch.save(model, 'model.pt')），可以直接加载：

model = torch.load('model.pt')

注意事项

1. 安全性：加载不受信任的文件可能导致恶意代码执行。建议使用 weights_only=True 限制仅加载张量。

   safe_data = torch.load('file.pt', weights_only=True)
   

2. 设备兼容性：如果模型保存时在 GPU，加载时需指定 map_location='cpu' 或目标设备。
3. 文件格式：支持 .pt 或 .pth 后缀，但需确保文件内容为 PyTorch 序列化数据。

错误处理

• 文件不存在时抛出 FileNotFoundError。
• 数据格式不匹配时抛出 pickle.UnpicklingError 或 RuntimeError。