浅谈pytorch 模型 .pt, .pth, .pkl的区别及模型保存方式

在PyTorch中，我们可以使用不同的文件格式来保存模型，包括.pt、.pth和.pkl。这些文件格式之间有一些区别，本文将对它们进行详细讲解，并提供两个示例说明。

.pt和.pth文件

.pt和.pth文件是PyTorch中最常用的模型保存格式。它们都是二进制文件，可以保存模型的参数、状态和结构。.pt文件通常用于保存单个模型，而.pth文件通常用于保存多个模型，例如在训练过程中保存的多个检查点。

以下是一个示例，展示如何将模型保存为.pt文件：

import torch
import torch.nn as nn

# Define model
model = nn.Linear(10, 1)

# Define input tensor
x = torch.randn(1, 10)

# Define output tensor
y = model(x)

# Save model
torch.save(model.state_dict(), 'model.pt')

在这个示例中，我们首先定义了一个线性模型model，它有10个输入和1个输出。接下来，我们定义了一个输入张量x，它的形状为(1, 10)。然后，我们将输入张量x应用于模型，得到输出张量y。最后，我们使用torch.save函数将模型的状态字典保存为model.pt文件。

以下是一个示例，展示如何将模型保存为.pth文件：

import torch
import torch.nn as nn

# Define model
model1 = nn.Linear(10, 1)
model2 = nn.Linear(10, 1)

# Define input tensor
x = torch.randn(1, 10)

# Define output tensor
y1 = model1(x)
y2 = model2(x)

# Save models
torch.save({
    'model1_state_dict': model1.state_dict(),
    'model2_state_dict': model2.state_dict()
}, 'models.pth')

在这个示例中，我们首先定义了两个线性模型model1和model2，它们都有10个输入和1个输出。接下来，我们定义了一个输入张量x，它的形状为(1, 10)。然后，我们将输入张量x分别应用于两个模型，得到输出张量y1和y2。最后，我们使用torch.save函数将两个模型的状态字典保存为models.pth文件。

.pkl文件

.pkl文件是Python中常用的序列化文件格式，可以保存任何Python对象，包括模型、数据和配置。.pkl文件通常用于保存整个模型，包括模型的参数、状态和结构。

以下是一个示例，展示如何将模型保存为.pkl文件：

import torch
import torch.nn as nn
import pickle

# Define model
model = nn.Linear(10, 1)

# Define input tensor
x = torch.randn(1, 10)

# Define output tensor
y = model(x)

# Save model
with open('model.pkl', 'wb') as f:
    pickle.dump(model, f)

在这个示例中，我们首先定义了一个线性模型model，它有10个输入和1个输出。接下来，我们定义了一个输入张量x，它的形状为(1, 10)。然后，我们将输入张量x应用于模型，得到输出张量y。最后，我们使用pickle.dump函数将整个模型保存为model.pkl文件。

总结

在本文中，我们详细讲解了PyTorch中的模型保存方式，包括.pt、.pth和.pkl文件，并提供了两个示例说明。.pt和.pth文件通常用于保存模型的参数和状态字典，而.pkl文件通常用于保存整个模型。