TinyNeuralNetwork
TinyNeuralNetwork是一个高效且易用的深度学习模型压缩框架,包含神经架构搜索、剪枝、量化、模型转换等功能。它已被应用于天猫精灵、海尔电视、优酷视频、人脸识别打卡机等设备的部署,为超过1000万台物联网设备提供了AI能力。
安装
Python版本 >= 3.8,PyTorch版本 >= 1.4(如果涉及量化感知训练,则PyTorch版本 >= 1.6,详情请参见此处)
# 安装TinyNeuralNetwork框架
git clone https://github.com/alibaba/TinyNeuralNetwork.git
cd TinyNeuralNetwork
python setup.py install
# 或者,您可以尝试使用以下一行命令
pip install git+https://github.com/alibaba/TinyNeuralNetwork.git
您也可以使用Docker构建
sudo docker build -t tinynn:pytorch1.9.0-cuda11.1 .
贡献
我们感谢您为改进我们的框架提供帮助。更多详情请参见此处。
基本模块
- 计算图捕获:TinyNeuralNetwork中的Graph Tracer捕获PyTorch算子的连接关系,实现自动化剪枝和模型量化。它还支持从PyTorch模型生成等效的模型描述文件(如models.py)。
- 依赖关系解析:修改一个算子通常会导致子图不匹配,即与其他依赖算子不匹配。TinyNeuralNetwork中的Graph Modifier自动处理子图内部和子图之间的不匹配,实现计算图的自动修改。
- 剪枝器:已实现OneShot(L1、L2、FPGM)、ADMM、NetAdapt、Gradual、End2End等剪枝算法,将逐步开放。
- 量化感知训练:TinyNeuralNetwork使用PyTorch的QAT作为后端(我们还支持模拟bfloat16训练),并通过自动化融合算子和量化计算图优化其可用性(官方实现需要用户手动实现,工作量巨大)。
- 模型转换:TinyNeuralNetwork支持将浮点和量化后的PyTorch模型转换为TFLite模型,实现端到端部署。
项目架构
路线图
- 2021年11月:具有自适应稀疏性的新剪枝器
- 2021年12月:Transformer模型压缩
引用
如果您在研究中发现本项目有用,请考虑引用:
@misc{tinynn,
title={TinyNeuralNetwork: An efficient deep learning model compression framework},
author={Ding, Huanghao and Pu, Jiachen and Hu, Conggang},
howpublished = {\url{https://github.com/alibaba/TinyNeuralNetwork}},
year={2021}
}
常见问题
由于PyTorch的高复杂性和频繁更新,我们无法通过自动化测试确保覆盖所有情况。 当您遇到问题时,可以查看常见问题,或通过下方二维码加入钉钉问答群。