【剪枝实战】使用VGGNet训练、稀疏训练、剪枝、微调等,剪枝出只有3M的模型

摘要

本次剪枝实战是基于下面这篇论文去复现的,主要是实现对BN层的γ/gamma进行剪枝操作,本文用到的代码和数据集都可以在我的资源中免费下载到。

相关论文:Learning Efficient Convolutional Networks through Network Slimming (ICCV 2017)

剪枝的原理

在一个卷积-BN-激活模块中,BN层可以实现通道的缩放。如下:

BN层的具体操作有两部分:

在归一化后会进行线性变换,那么当系数gamma很小时候,对应的激活(Zout)会相应很小。这些响应很小的输出可以裁剪掉,这样就实现了bn层的通道剪枝。

剪枝的过程

第一步、使用VGGNet训练模型。保存训练结果,方便将来的比对!

第二步、在BN层网络中加入稀疏因子,训练模型。

第三步、剪枝操作。

第四步、fine-tune模型,提高模型的ACC。

接下来,我们一起实现对VGGNet的剪枝。

一、项目结构

Slimming_Demo

├─checkpoints

│ ├─vgg

│ ├─vgg_pruned

│ └─vgg_sp

├─data

│ ├─train

│ │ ├─Black-grass

│ │ ├─Charlock

│ │ ├─Cleavers

│ │ ├─Common Chickweed

│ │ ├─Common wheat

│ │ ├─Fat Hen

│ │ ├─Loose Silky-bent

│ │ ├─Maize

│ │ ├─Scentless Mayweed

│ │ ├─Shepherds Purse

│ │ ├─Small-flowered Cranesbill

│ │ └─Sugar beet

│ └─val

│ ├─Black-grass

│ ├─Charlock

│ ├─Cleavers

│ ├─Common Chickweed

│ ├─Common wheat

│ ├─Fat Hen

│ ├─Loose Silky-bent

│ ├─Maize

│ ├─Scentless Mayweed

│ └─Shepherds Purse

├─vgg.py

├─train.py

├─train_sp.py

├─prune.py

└─train_prune.py

train.py:训练脚本,训练VGGNet原始模型 vgg.py:模型脚本 train_sp.py:稀疏训练脚本。 prune.py:模型剪枝脚本。 train_prune.py:微调模型脚本。

二、正常训练

首先,我们进行正常训练出一个模型,用来和剪枝后的模型进行对比。原始模型有67.8M,体积非常大。

运行train.py,正常训练结果保存为checkpoints/vgg/best.pth

通过Tensorboard工具,我们观察到,此时的BN层的gamma系数并不稀疏,很难判断出各个通道的重要性,不利于剪枝!

三、稀疏训练

运行train_sp.py,稀疏化训练结果保存为checkpoints/vgg_sp/best.pth

使用tensorboard查看runs/路径下保存的log,查看稀疏化训练结果。

可以看到通过稀疏化训练,==此时的gamma系数已经呈现趋近0的分布,==我们可以对gamma接近0的channel进行剪枝!

核心代码

四、剪枝

有了稀疏化训练的模型,我们便可以通过读取checkpoints/vgg_sp/best.pth,根据gamma系数来判断网络中各个通道的重要性,剪去不重要的通道。

运行prune.py,剪枝的结果保存为checkpoints/vgg_pruned/best.pth

查看剪枝后的VGG模型,只有3.58M,体积非常小,剪枝瘦身效果非常明显!

五、微调

对模型剪枝后,为了恢复模型的性能,我们还需要对剪枝的模型进行Finetune

运行train_prune.py

checkpoints/vgg_pruned/best.pth进行微调,结果生成在checkpoints/vgg_pruned_finetune/best.pth

六、总结

我们通过复现Learning Efficient Convolutional Networks through Network Slimming (ICCV 2017)论文里提到的网络大瘦身剪枝算法,完成对VGG模型进行剪枝。

相关推荐
AI极客菌30 分钟前
Controlnet作者新作IC-light V2:基于FLUX训练,支持处理风格化图像,细节远高于SD1.5。
人工智能·计算机视觉·ai作画·stable diffusion·aigc·flux·人工智能作画
阿_旭32 分钟前
一文读懂| 自注意力与交叉注意力机制在计算机视觉中作用与基本原理
人工智能·深度学习·计算机视觉·cross-attention·self-attention
王哈哈^_^38 分钟前
【数据集】【YOLO】【目标检测】交通事故识别数据集 8939 张,YOLO道路事故目标检测实战训练教程!
前端·人工智能·深度学习·yolo·目标检测·计算机视觉·pyqt
Power20246662 小时前
NLP论文速读|LongReward:基于AI反馈来提升长上下文大语言模型
人工智能·深度学习·机器学习·自然语言处理·nlp
数据猎手小k2 小时前
AIDOVECL数据集:包含超过15000张AI生成的车辆图像数据集,目的解决旨在解决眼水平分类和定位问题。
人工智能·分类·数据挖掘
好奇龙猫2 小时前
【学习AI-相关路程-mnist手写数字分类-win-硬件:windows-自我学习AI-实验步骤-全连接神经网络(BPnetwork)-操作流程(3) 】
人工智能·算法
沉下心来学鲁班2 小时前
复现LLM:带你从零认识语言模型
人工智能·语言模型
数据猎手小k2 小时前
AndroidLab:一个系统化的Android代理框架,包含操作环境和可复现的基准测试,支持大型语言模型和多模态模型。
android·人工智能·机器学习·语言模型
YRr YRr2 小时前
深度学习:循环神经网络(RNN)详解
人工智能·rnn·深度学习
sp_fyf_20242 小时前
计算机前沿技术-人工智能算法-大语言模型-最新研究进展-2024-11-01
人工智能·深度学习·神经网络·算法·机器学习·语言模型·数据挖掘