pytorch读取.pth文件

最新推荐文章于 2025-02-19 17:39:49 发布

蒲建建

最新推荐文章于 2025-02-19 17:39:49 发布

阅读量2.8w

点赞数 11

分类专栏：计算机视觉应用

原文链接：https://www.jianshu.com/p/946d1ce9c149

版权

计算机视觉应用专栏收录该内容

26 篇文章

订阅专栏

本文详细介绍了PyTorch中.pth文件的结构，它通过有序字典保存模型参数，每个元素为Parameter类型。讲解了torch.save()和torch.load()的用法，以及如何加载预训练模型的部分参数。在恢复训练或测试时，可以加载state_dict和optimizer的状态。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

1.pth文件中保存的是什么

import torch
state_dict = torch.load("resnet18.pth")
print(type(state_dict))

---------------
<class 'collections.OrderedDict'>

如上打印输出所示，pth文件通过有序字典来保持模型参数。有序字典与常规字典一样，但是在排序操作方面有一些额外的功能。常规的dict是无序的，OrderedDict能够比dict更好地处理频繁的重新排序操作。
OrderedDict有一个方法popitem(last=True)用于有序字典的popitem()方法返回并删除一个(键，值)对。如果last为真，则按LIFO顺序返回对;如果为假，则按FIFO顺序返回对。
OrderedDict还有一个方法move_to_end(key,last=True),将现有的键移动到有序字典的两端。如果last为真，则将项目移动到右端(默认);如果last为假，则移动到开头。

import torch
state_dict = torch.load("resnet18.pth")
print(type(state_dict))

for i in state_dict:
    print(i)
    print(type(state_dict[i]))
    print("aa:",state_dict[i].data.size())
    print("bb:",state_dict[i].requires_grad)
    break

------------------------------
<class 'collections.OrderedDict'>
conv1.weight
<class 'torch.nn.parameter.Parameter'>
aa: torch.Size([64, 3, 7, 7])
bb: True

如上打印所示，有序字典state_dict中每个元素都是Parameter参数，该参数是一种特殊的张量，包含data和requires_grad两个方法。其中data字段保存的是模型参数，requires_grad字段表示当前参数是否需要进行反向传播。

2.torch.save()

先建立一个字典，保存三个参数：调用torch.save(),即可保存对应的pth文件。需要注意的是若模型是由nn.Moudle类继承的模型，保存pth文件时，state_dict参数需要由model.state_dict指定。

state_dict = {‘net':model.state_dict(), 'optimizer':optimizer.state_dict(), 'epoch':epoch}
torch.save(state_dict , dir)

--------------------------------
torch.save(model.state_dict,dir)

3.torch.load()

当你想恢复某一阶段的训练（或者进行测试）时，那么就可以读取之前保存的网络模型参数等。

checkpoint = torch.load(dir)
model.load_state_dict(checkpoint['net'])
optimizer.load_state_dict(checkpoint['optimizer'])
start_epoch = checkpoint['epoch'] + 1

pytorch加载预训练模型部分参数

resnet = models.resnet50(pretrained=True)
new_state_dict = resnet.state_dict()
dd = net.state_dict()  #net是自己定义的含有resnet backbone的模型
for k in new_state_dict.keys():
    print(k)
    if k in dd.keys() and not k.startswith('fc'):  #不使用全连接的参数
        print('yes')
        dd[k] = new_state_dict[k]
net.load_state_dict(dd)

更加全面参考：https://blog.csdn.net/weixin_41519463/article/details/103205665