前言:前面有专门的讲解关于如何深入查询模型的参数信息,可以参考这篇文章:
本次来解析一下我们通常保存的模型文件 .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.