深度学习中的.pt和.pth文件的区别

什么是.pt和.pth文件？

首先，.pt和.pth文件都是用于保存PyTorch模型的文件格式。它们可以用于保存模型的权重、结构以及其他相关信息。

.pt文件

.pt文件是PyTorch模型保存的通用格式。通常，.pt文件可以保存以下内容：

1. 模型的权重（weights）：这是最常见的用法，只保存模型的参数。
2. 完整的模型（full model）：包括模型的架构和权重。在这种情况下，可以通过简单的一行代码来加载整个模型。

保存模型权重代码

import torch

# 假设有一个训练好的模型 model
torch.save(model.state_dict(), 'model_weights.pt')

保存完整模型代码

import torch

# 假设有一个训练好的模型 model
torch.save(model, 'full_model.pt')

加载模型权重底阿妈

import torch

# 假设有一个模型类 ModelClass
model = ModelClass()
model.load_state_dict(torch.load('model_weights.pt'))

加载完整模型代码

import torch

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

.pth文件

.pth文件在功能上与.pt文件类似，但在实际使用中通常有一些惯例上的区别。具体来说，.pth文件通常用于以下两种情况：

1. 保存和加载模型权重：与.pt文件类似，但在社区中更常用于保存训练好的模型权重。
2. 保存和加载优化器的状态：在训练过程中，我们可能希望保存优化器的状态，以便在训练中断后能够从中断点继续。

保存模型权重代码：

import torch

# 假设有一个训练好的模型 model
torch.save(model.state_dict(), 'model_weights.pth')

保存优化器状态代码：

import torch

# 假设有一个训练好的优化器 optimizer
torch.save(optimizer.state_dict(), 'optimizer_state.pth')

加载模型权重的代码：

import torch

# 假设有一个模型类 ModelClass
model = ModelClass()
model.load_state_dict(torch.load('model_weights.pth'))

加载优化器状态的代码：

import torch

# 假设有一个优化器 optimizer
optimizer.load_state_dict(torch.load('optimizer_state.pth'))