pytorch 模型保存和加载

最新推荐文章于 2025-07-01 12:29:59 发布

原创最新推荐文章于 2025-07-01 12:29:59 发布 · 192 阅读

1 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch #人工智能 #python

深度学习专栏收录该内容

8 篇文章

订阅专栏

本文介绍了如何在PyTorch中保存和加载模型，包括使用pkl/pt/pth格式，保存整个模型的结构和参数，以及只保存和加载模型权重。加载时需注意先构建模型结构再加载权重。

pytorch 模型保存和加载

1. 保存
2. 加载

1. 保存

格式：pkl/pt/pth

model = models.resnet50()
model.cuda()
save_dir = 'resnet50.pth' #保存路径

# 保存整个模型（结构+参数）
torch.save(model, save_dir)

# 保存模型权重(不要忘记加())
torch.save(model.state_dict(), save_dir)

2. 加载

# 读取整个模型
model_load = torch.load(save_dir)
model_load.cuda()

# 读取模型权重（需先加载模型）
model_load = models.resnet50()
# model_load = MyModel()
model_load.load_state_dict(torch.load(save_dir))
model_load.cuda()