pt模型剪裁

AI小笔记

已于 2022-01-21 14:14:34 修改

阅读量1.8k

点赞数

CC 4.0 BY-SA版权

分类专栏：深度学习文章标签： pytorch 深度学习

于 2021-10-08 14:28:17 首次发布

本文链接：https://blog.youkuaiyun.com/JulyLi2019/article/details/120650969

深度学习专栏收录该内容

30 篇文章

订阅专栏

本文介绍了如何加载.pt模型文件，查询并删除无意义的键值对，如训练过程中的统计数据和优化器参数，以减小模型体积。通过删除不必要的键值对，例如'optimizer'、'training_results'等，可以有效地优化模型存储。最后展示了如何保存精简后的模型权重文件。

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

.pt文件是模型权重，主要包含训练过程中保留的键值对信息
以自己训练的yolov5s目标检测权重test.pt为例：

1.模型文件加载

pt = r'weights/test.pt'
state = torch.load(pt)

2.模型键值对查询

for k in state.keys():
print(k)

epoch： 113（对后续推理无意义可删除）
best_fitness： [0.85533715] （对后续推理无意义可删除）
training_results： 0/599 2.47G 0.05426… （对后续推理无意义可删除）
model: 模型权重
ema：模型权重（重复可删除）
updates： 23620 （对后续推理无意义可删除）
optimizer：优化器的一些参数（只用作推理可删除）
wandb_id：可视化序列值（可删除）

3.删除无效键值对（del）

del state['optimizer']['state']             
del state['training_results']               
del state['ema']							  
del state['updates']

4.查询键值类型（type）

print(type(state['model']))     				<class 'models.yolo.Model'>
print(type(state['updates']))					<class 'int'>
print(type(['wandb_id']))						<class 'list'>
print(type(state['optimizer']))				<class 'dict'>

5.模型文件保存

state：加载的模型权重
filename：存储的文件名

torch.save(state, filename)

去掉无用的信息后，pt模型体积会减小很多：
在这里插入图片描述

6.实例代码

以yolov5模型为例：

import torch

pt = r'weights/best.pt'
# 加载
state = torch.load(pt)

# 键值对
for k in state.keys():
    print(k)
del state['epoch']
del state['best_fitness']
del state['training_results']
del state['ema']
del state['updates']
del state['optimizer']
del state['wandb_id']

# 保存
torch.save(state, 'weights/best_cut.pt')

for k in state.keys():
    print(k)