pytorch中保存的模型文件.pth深入解析

前言:前面有专门的讲解关于如何深入查询模型的参数信息,可以参考这篇文章:

沈鹏燕:pytorch教程之nn.Module类详解zhuanlan.zhihu.com图标

本次来解析一下我们通常保存的模型文件 .pth 文件到底内部是什么?

一、.pth 文件详解

在pytorch进行模型保存的时候,一般有两种保存方式,一种是保存整个模型,另一种是只保存模型的参数。

torch.save(model.state_dict(), "my_model.pth") # 只保存模型的参数

torch.save(model, "my_model.pth") # 保存整个模型
保存的模型参数实际上一个字典类型,通过key-value的形式来存储模型的所有参数,本文以自己在实践过程中使用的一个.pth文件为例来说明,使用的是整个模型。

1.1 .pth 文件基本信息的查看

import torch

pthfile = r'F:/GNN/graph-rcnn/graph-rcnn/datasets/sg_baseline_ckpt.pth' #faster_rcnn_ckpt.pth
net = torch.load(pthfile,map_location=torch.device('cpu')) # 由于模型原本是用GPU保存的,但我这台电脑上没有GPU,需要转化到CPU上

# print(type(net)) # 类型是 dict
# print(len(net)) # 长度为 4,即存在四个 key-value 键值对

# for k in net.keys():
# print(k) # 查看四个键,分别是 model,optimizer,scheduler,iteration
1.2 模型的四个键值分别详解

(1)net["model"] 详解

# print(net["model"]) # 返回的是一个OrderedDict 对象
for key,value in net["model"].items():
print(key,value.size(),sep=" ")
'''运行结果如下:
module.backbone.body.stem.conv1.weight torch.Size([64, 3, 7, 7])
module.backbone.body.stem.bn1.weight torch.Size([64])
module.backbone.body.stem.bn1.bias torch.Size([64])
module.backbone.body.stem.bn1.running_mean torch.Size([64])
module.backbone.body.stem.bn1.running_var torch.Size([64])
module.backbone.body.layer1.0.downsample.0.weight torch.Size([256, 64, 1, 1])
module.backbone.body.layer1.0.downsample.1.weight torch.Size([256])
.
.
.
module.backbone.body.layer3.22.bn3.weight torch.Size([1024])
module.backbone.body.layer3.22.bn3.bias torch.Size([1024])
module.backbone.body.layer3.22.bn3.running_mean torch.Size([1024])
module.backbone.body.layer3.22.bn3.running_var torch.Size([1024])
.
.
.
module.rpn.head.conv.bias torch.Size([1024])
module.rpn.head.cls_logits.weight torch.Size([15, 1024, 1, 1])
module.rpn.head.cls_logits.bias torch.Size([15])
module.rpn.head.bbox_pred.weight torch.Size([60, 1024, 1, 1])
.
.
.
module.roi_heads.box.feature_extractor.head.layer4.0.bn2.running_var torch.Size([512])
module.roi_heads.box.feature_extractor.head.layer4.0.conv3.weight torch.Size([2048, 512, 1, 1])
module.roi_heads.box.feature_extractor.head.layer4.0.bn3.weight torch.Size([2048])
.
.
.
module.roi_heads.relation.predictor.cls_score.weight torch.Size([51, 2048])
module.roi_heads.relation.predictor.cls_score.bias torch.Size([51])

'''
总结:键model所对应的值是一个OrderedDict,而这个OrderedDict字典里面又存储着所有的每一层的参数名称以及对应的参数值。

需要注意的是,这里参数名称之所以很长,如:

module.backbone.body.stem.conv1.weight
是因为搭建网络结构的时候采用了组件式的设计,即整个模型里面构造了一个backbone的容器组件,backbone里面又构造了一个body容器组件,body里面又构造了一个stem容器,stem里面的第一个卷积层的权重。

(2)net["optimizer"]详解

# print(net["optimizer"]) # 返回的是一个一般的字典 Dict 对象
for key,value in net["optimizer"].items():
print(key,type(value),sep=" ")
'''运行结果为:
state <class 'dict'>
param_groups <class 'list'>
'''
'''
发现这个这个字典只有两个key,一个是state,一个是param_groups
其中state所对应的值又是一个字典类型,
param_groups对应的值是一个列表
'''
继续往下查看得到

先看一下net["optimizer"]["param_groups"] 这个列表里面放了一下啥:

groups=net["optimizer"]["param_groups"]
print(groups)
print(len(groups)) # 返回115.即在这个模型中,共有115组

'''
[{'lr': 5.000000000000001e-05, 'weight_decay': 0.0005, 'momentum': 0.9, 'dampening': 0, 'nesterov': False, 'initial_lr': 0.005, 'params': [140566644061240]},
{'lr': 5.000000000000001e-05, 'weight_decay': 0.0005, 'momentum': 0.9, 'dampening': 0, 'nesterov': False, 'initial_lr': 0.005, 'params': [140566644061960]},
{'lr': 5.000000000000001e-05, 'weight_decay': 0.0005, 'momentum': 0.9, 'dampening': 0, 'nesterov': False, 'initial_lr': 0.005, 'params': [140566644062248]},
{'lr': 5.000000000000001e-05, 'weight_decay': 0.0005, 'momentum': 0.9, 'dampening': 0, 'nesterov': False, 'initial_lr': 0.005, 'params': [140566644077336]},
.
.
.
{'lr': 5.000000000000001e-05, 'weight_decay': 0.0005, 'momentum': 0.9, 'dampening': 0, 'nesterov': False, 'initial_lr': 0.005, 'params': [140566644061960]},
{'lr': 5.000000000000001e-05, 'weight_decay': 0.0005, 'momentum': 0.9, 'dampening': 0, 'nesterov': False, 'initial_lr': 0.005, 'params': [140566103171936]},
{'lr': 5.000000000000001e-05, 'weight_decay': 0.0005, 'momentum': 0.9, 'dampening': 0, 'nesterov': False, 'initial_lr': 0.005, 'params': [140566103172008]}
]
'''
这个列表的长度为115,每一个元素又是一个字典。

再看一下net["optimizer"]["states"] 这个字典里面放了啥:

state=net["optimizer"]["state"]
print(len(state)) # 返回115.即在这个模型中,state共有115组

for key,value in state.items():
print(key,type(value),sep=" ")
'''
140566644061240 <class 'dict'>
140566644061960 <class 'dict'>
140566644062248 <class 'dict'>
140566644077336 <class 'dict'>
.
.
.
140566103171936 <class 'dict'>
140566103172008 <class 'dict'>
'''
这个字典的长度是115,而且和前面的param_groups有着对应关系,每一个元素的键值就是param_groups中每一个元素的params。

继续往深一层看:

print(type(state[140566644061240])) # 他又是一个字典
for key,value in state[140566644061240].items():
print(key,value.size(),sep=" ")

'''
<class 'dict'>
momentum_buffer torch.Size([512, 256, 1, 1])
'''
(3)net["scheduler"] 详解

scheduler=net["scheduler"] # 返回的依然是一个字典
print(len(scheduler)) # 字典的长度为 7
print(scheduler)
'''
{'milestones': (70000, 90000),
'gamma': 0.1,
'warmup_factor': 0.3333333333333333,
'warmup_iters': 500,
'warmup_method': 'linear',
'base_lrs': [0.005, 0.005, 0.005, 0.01, ......, 0.005, 0.005, 0.005, 0.005, 0.01],
'last_epoch': 99999}
'''
继续看一下这个base_lrs的信息

print(len(scheduler["base_lrs"])) # 返回115,→115个数组成的一个列表
(4)net["iteration"] 详解

print(net["iteration"]) # 返回 9999 ,它是一个具体的数字


二、关于.pth 文件的总结

它是一个包含 四组 “key-value”的字典,类型分别如下:

其中

(1)net["model"] 就相当于是 前面文章中说到的 net.state_dict() 返回的那个字典;

(2)net["optimizer"] 就相当于是 前面文章中说到的 optimizer.state_dict() 返回的那个字典;

发布于 2019-09-30

文章被以下专栏收录