【几分钟】快速熟悉torch.save()、torch.load()、torch.nn.Module.load_state_dict()

【几分钟】快速熟悉torch.save()、torch.load()、torch.nn.Module.load_state_dict()

🌳引言🌳

在PyTorch中，模型训练完成后通常需要保存以便后续使用或进行进一步的训练。PyTorch提供了几种方法来实现模型的保存和加载，其中torch.save(), torch.load()和torch.nn.Module.load_state_dict()是最常用的函数。本文将用几分钟的时间带您快速熟悉这三个函数的使用方法和注意事项。

🌳torch.save()详解🌳

torch.save()函数用于保存模型的状态或整个模型。其用法如下：

torch.save(obj, f)

• obj: 要保存的对象，可以是模型的状态字典、整个模型等。
• f: 保存文件的路径。

当有保存模型的需求时，通常推荐只保存模型的参数（即状态字典），而不是整个模型实例。这样可以避免保存模型定义时的额外信息，比如优化器的状态等，保存模型的示例如下：

# 保存模型的状态字典
torch.save(model.state_dict(), 'model_state_dict.pth')

# 如果需要保存整个模型，可以这样做，但通常不推荐
torch.save(model, 'model.pth')

🌳torch.load()详解🌳

torch.load()函数用于加载之前保存的模型或状态字典。其用法如下：

torch.load(f, map_location=None)

• f: 加载文件的路径。
• map_location: 指定加载模型到哪个设备上，比如CPU或特定的GPU。

加载模型时，需要根据保存时的方式选择加载整个模型还是仅加载状态字典。

# 加载状态字典
state_dict = torch.load('model_state_dict.pth')

# 加载整个模型（如果之前是这样保存的）
model = torch.load('model.pth')

🌳torch.nn.Module.load_state_dict()详解🌳

torch.nn.Module.load_state_dict()是PyTorch模型类（继承自torch.nn.Module）的一个方法，用于加载状态字典。其用法如下：

model.load_state_dict(state_dict, strict=True)

• state_dict: 要加载的状态字典。
• strict: 是否严格检查加载的状态字典与模型当前的状态字典是否完全匹配。默认为True。

使用load_state_dict()加载状态字典时，需要先实例化模型类，然后调用此方法加载之前保存的状态。

# 实例化模型类
model = MyModel()

# 加载状态字典
model.load_state_dict(torch.load('model_state_dict.pth'))

🌳保存并加载模型的几种方式🌳

1. 仅保存和加载状态字典

   这是推荐的方式，因为它只保存和加载模型的参数，不包含其他不必要的信息。

# 保存
torch.save(model.state_dict(), 'model_state_dict.pth')

# 加载
model = MyModel()
model.load_state_dict(torch.load('model_state_dict.pth'))

2. 保存和加载整个模型

   这种方式会保存模型的所有信息，包括参数、优化器状态等。但这种方式不够灵活，通常不推荐。

# 保存
torch.save(model, 'model.pth')

# 加载
model = torch.load('model.pth')

🌳总结🌳

在PyTorch中，模型的保存和加载主要通过torch.save(), torch.load()和torch.nn.Module.load_state_dict()实现。推荐的做法是只保存和加载模型的状态字典，这样更加灵活且只包含模型的核心信息。在加载模型时，需要先实例化模型类，然后使用load_state_dict()方法加载状态字典。

🌳结尾🌳

亲爱的读者，首先感谢您抽出宝贵的时间来阅读我们的博客。我们真诚地欢迎您留下评论和意见💬。
俗话说，当局者迷，旁观者清。您的客观视角对于我们发现博文的不足、提升内容质量起着不可替代的作用。
如果博文给您带来了些许帮助，那么，希望您能为我们点个免费的赞👍👍/收藏👇👇，您的支持和鼓励👏👏是我们持续创作✍️✍️的动力。
我们会持续努力创作✍️✍️，并不断优化博文质量👨‍💻👨‍💻，只为给您带来更佳的阅读体验。
如果您有任何疑问或建议，请随时在评论区留言，我们将竭诚为你解答~
愿我们共同成长🌱🌳，共享智慧的果实🍎🍏！

万分感谢🙏🙏您的点赞👍👍、收藏⭐🌟、评论💬🗯️、关注❤️💚~