Pytorch 什么是Pytorch中.ckpt和.pth文件的区别

在本文中，我们将介绍Pytorch中.ckpt和.pth文件的区别。Pytorch是一个广泛使用的深度学习框架，.ckpt和.pth文件是在该框架中保存模型权重的两种常见的文件格式。我们将详细解释这两种文件格式，并举例说明其用法和区别。

.ckpt文件

.ckpt文件是Pytorch中的一种标准模型检查点文件格式，通常用于保存模型的状态字典。这个文件格式是为了方便在训练过程中保存和恢复模型所用的。保存为.ckpt文件的模型可以很容易地在同一个框架内加载并继续训练，而无需重新从头开始。

.ckpt文件通常包含了模型中所有的权重参数和优化器状态等信息，以便能够完全重现模型的当前状态。这对于长时间训练的模型、模型参数的共享和复现以及断点续训（checkpointing）非常有用。

下面是一个示例代码，展示如何保存和加载.ckpt文件：

import torch
import torch.nn as nn
import torch.optim as optim

# 创建模型
model = nn.Linear(10, 2)
optimizer = optim.SGD(model.parameters(), lr=0.01)

# 保存模型状态字典
torch.save({
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
}, 'model.ckpt')

# 加载模型状态字典
checkpoint = torch.load('model.ckpt')
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])

如上所示，首先我们创建了一个简单的模型和优化器，然后使用torch.save()函数将模型的状态字典和优化器的状态字典保存为.ckpt文件。然后，我们使用torch.load()函数加载.ckpt文件，并使用load_state_dict()方法将状态字典恢复到模型和优化器中。

.pth文件

.pth文件是Pytorch中另一种常见的模型文件格式，通常用于保存整个模型，包括模型的架构、权重和其他相关信息。.pth文件比.ckpt文件更为完整，可以方便地在其他框架或不同的设备上加载和使用。

下面是一个示例代码，展示如何保存和加载.pth文件：

import torch
import torch.nn as nn
import torch.optim as optim

# 创建模型
model = nn.Linear(10, 2)
optimizer = optim.SGD(model.parameters(), lr=0.01)

# 保存整个模型
torch.save(model, 'model.pth')

# 加载整个模型
model = torch.load('model.pth')

如上所示，我们可以直接使用torch.save()函数将整个模型保存为.pth文件，加载时也可以使用torch.load()函数直接加载整个模型。