pth.tar的保存和读取

yrhzmu2024-01-04 15:21

一、简介

在PyTorch中，.pt、.pth和.pth.tar都是保存训练好的模型的文件格式。主要区别在于：

.pt是PyTorch1.6及以上版本中引入的保存格式，可以保存整个模型，包括模型结构、模型参数以及优化器状态等信息，是一个二进制文件。
.pth是PyTorch旧版本中使用的模型文件格式，只保存了模型参数，没有保存模型和其他相关信息，是一个二进制文件。
.pth.tar包括.pth文件以及其他信息，比如模型结构、优化器状态、超参数信息。

二、保存

使用torch.save进行保存，保存时传入保存的状态，名称

python 复制代码

def save_checkpoint(state, is_best, filename='checkpoint.pth.tar'):
    torch.save(state, filename)
    if is_best:
        shutil.copyfile(filename, 'model_best.pth.tar')

三、读取

1、通过torch.load()函数加载

python 复制代码

checkpoint_path = "/home/user/msh/Project/SimCLR-master_old/runs/Jan03_19-04-59_user-X10DRi/checkpoint_0100.pth.tar"
checkpoint = torch.load(checkpoint_path)
print(checkpoint.keys())

运行结果如下:

python 复制代码

dict_keys(['epoch', 'arch', 'state_dict', 'optimizer'])

2、epoch存放的是训练的轮次，arch存放的是模型的名称，optimizer存放是优化器具体的参数，

python 复制代码

epoch = checkpoint['epoch']
print(epoch)
arch = checkpoint['arch']
print(arch)
optimizer = checkpoint['optimizer']

运行结果：

python 复制代码

100
resnet18

3、state_dict.keys()存放的是模型每一层结构的名称

python 复制代码

state_dict = checkpoint['state_dict']
print(state_dict.keys())

4、使用：先初始化模型，创建一个对象，然后使用load_state_dict()函数加载参数

python 复制代码

model = ResNetSimCLR(arch,160)
model.load_state_dict(state_dict)