.pt文件
.pt
文件是 PyTorch 的模型文件,用于保存模型的权重、结构或者两者兼有。这些文件可以包含以下内容之一:
- 仅模型权重(state_dict):保存模型的参数(权重和偏置)。
- 完整模型:包括模型的结构和权重(通常是通过 JIT 编译保存的 TorchScript 模型)。
- 其它对象:如优化器状态、训练状态等。
查看和加载 .pt
文件的内容可以通过以下方法进行:
1. 仅模型权重(state_dict)
如果 .pt
文件保存的是模型的权重,可以使用 torch.load
加载它,并查看包含的键和值:
import torch
# 假设文件名为 'model_weights.pt'
state_dict = torch.load('model_weights.pt')
# 打印 state_dict 中的键
print(state_dict.keys())
# 查看某个具体键的内容
print(state_dict['some_layer.weight'])
2. 完整模型(TorchScript 模型)
如果 .pt
文件保存的是通过 TorchScript 编译的完整模型,可以使用 torch.jit.load
加载它,并查看模型结构:
import torch
# 假设文件名为 'model_scripted.pt'
model = torch.jit.load('model_scripted.pt')
# 打印模型结构
print(model)
# 进行前向传递
dummy_input = torch.randn(1, 3, 224, 224) # 假设输入是图像
output = model(dummy_input)
print(output)
3. 其它对象
有时 .pt
文件可能包含多个对象,例如模型和优化器的状态。可以使用 torch.load
加载并查看:
import torch
# 假设文件名为 'checkpoint.pt'
checkpoint = torch.load('checkpoint.pt')
# 打印 checkpoint 中的键
print(checkpoint.keys())
# 加载模型权重和优化器状态
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
示例:查看 .pt
文件内容
假设你有一个 .pt
文件,名为 example_model.pt
,你可以使用以下代码加载并查看其内容:
import torch
# 加载 .pt 文件
content = torch.load('example_model.pt')
# 检查内容类型
if isinstance(content, dict):
# 打印字典中的键
print("Keys in the checkpoint:", content.keys())
# 假设它包含模型的 state_dict
if 'model_state_dict' in content:
model_state_dict = content['model_state_dict']
print("Model state dict keys:", model_state_dict.keys())
# 假设它包含优化器的 state_dict
if 'optimizer_state_dict' in content:
optimizer_state_dict = content['optimizer_state_dict']
print("Optimizer state dict keys:", optimizer_state_dict.keys())
else:
# 如果内容是一个模型
print(content)
总结
通过上述方法,可以查看和理解 .pt
文件的内容,并根据需要加载和使用这些内容。通常情况下,文件中会包含模型的权重或完整模型,以及其它训练相关的信息,这些都可以通过 torch.load
或 torch.jit.load
进行加载和查看。
.pth和.checkpoint文件
.pth
和 .checkpoint
文件与 .pt
文件类似,通常用于保存和加载 PyTorch 模型及其相关信息。具体来说:
.pth
文件通常用于保存和加载模型的权重。.checkpoint
文件通常用于保存训练检查点,包括模型的权重、优化器状态、训练进度等。
查看和加载 .pth
文件
.pth
文件通常保存模型的权重(state_dict)。你可以使用 torch.load
加载这些权重,然后将其加载到模型中:
import torch
import torch.nn as nn
# 定义你的模型结构
class MyModel(nn.Module):
def __init__(self):
super(MyModel, self).__init__()
self.fc = nn.Linear(10, 1)
def forward(self, x):
return self.fc(x)
# 加载模型权重
model = MyModel()
model.load_state_dict(torch.load('model.pth'))
# 现在可以使用模型进行推理
model.eval()
dummy_input = torch.randn(1, 10)
output = model(dummy_input)
print(output)
查看和加载 .checkpoint
文件
.checkpoint
文件通常包含更多的信息,比如模型权重、优化器状态、训练进度等。你可以使用 torch.load
加载它们,并提取需要的信息:
import torch
import torch.nn as nn
import torch.optim as optim
# 定义你的模型结构
class MyModel(nn.Module):
def __init__(self):
super(MyModel, self).__init__()
self.fc = nn.Linear(10, 1)
def forward(self, x):
return self.fc(x)
# 加载检查点
checkpoint = torch.load('model.checkpoint')
# 创建模型和优化器
model = MyModel()
optimizer = optim.SGD(model.parameters(), lr=0.001)
# 加载模型权重和优化器状态
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
# 加载训练进度
epoch = checkpoint['epoch']
loss = checkpoint['loss']
# 打印检查点中的信息
print(f"Checkpoint loaded: Epoch {epoch}, Loss {loss}")
# 继续训练或推理
model.eval()
dummy_input = torch.randn(1, 10)
output = model(dummy_input)
print(output)
示例:保存和加载 .checkpoint
文件
保存一个训练检查点的示例代码:
import torch
import torch.nn as nn
import torch.optim as optim
# 定义模型结构
class MyModel(nn.Module):
def __init__(self):
super(MyModel, self).__init__()
self.fc = nn.Linear(10, 1)
def forward(self, x):
return self.fc(x)
# 创建模型和优化器
model = MyModel()
optimizer = optim.SGD(model.parameters(), lr=0.001)
# 模拟训练过程
num_epochs = 5
for epoch in range(num_epochs):
# 假设我们计算了一些损失
loss = torch.tensor(epoch * 0.1)
# 保存检查点
checkpoint = {
'epoch': epoch,
'model_state_dict': model.load_state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'loss': loss,
}
torch.save(checkpoint, f'model_epoch_{epoch}.checkpoint')
print(f"Checkpoint saved for epoch {epoch}")
总结
.pth
文件通常用于保存模型的权重,可以使用torch.load
加载并应用到模型中。.checkpoint
文件通常用于保存训练检查点,包含模型权重、优化器状态和训练进度等信息。
通过这些示例和解释,你可以查看和加载 .pth
和 .checkpoint
文件的内容,并在需要时恢复模型和训练状态。