【最佳实践】.pth还是.tar？

最新推荐文章于 2024-09-04 09:55:21 发布

Sailist

最新推荐文章于 2024-09-04 09:55:21 发布

阅读量1.2w

点赞数 19

分类专栏： Pytorch

本文链接：https://blog.csdn.net/sailist/article/details/103862438

版权

Pytorch 专栏收录该内容

6 篇文章

订阅专栏

本文详细介绍了PyTorch中模型保存与加载的最佳实践，包括仅保存模型用于推断时使用.pth或.pt，以及断点保存使用.tar，并通过load_state_dict()方法加载。文章还解释了.pth.tar文件的使用场景及后缀名的选择依据。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

pytorch的官方教程里提供了相关说明：

只保存模型用于以后的推断的话使用.pth或.pt，这样可以直接加载模型

A common PyTorch convention is to save models using either a .pt or .pth file extension.

torch.save(model, "model.pth") # or .pt
model = torch.load("model.pth")

断点保存的话则使用.tar，加载的时候模型需要使用load_state_dict()方法

To save multiple components, organize them in a dictionary and use torch.save() to serialize the dictionary. A common PyTorch convention is to save these checkpoints using the .tar file extension.

torch.save({
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'loss': loss,
            ...
            }, "checkpoint.tar")

...

checkpoint = torch.load("checkpoint.tar")
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epoch = checkpoint['epoch']
loss = checkpoint['loss']