TinyNeuralNetwork

简体中文

TinyNeuralNetwork 是一个高效且易于使用的深度学习模型压缩框架，包含神经架构搜索、剪枝、量化、模型转换等功能。该框架已被应用于天猫精灵、海尔电视、优酷视频、人脸识别签到机等设备的部署中，为超过1000万台物联网设备赋予了AI能力。

安装

Python >= 3.9，PyTorch >= 1.10

# 安装 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 .

贡献

我们非常感谢您对本框架改进的支持。更多详情请参阅 CONTRIBUTING.md。

基础模块

• 计算图捕获：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 模型，以便进行端到端部署。

项目架构

• examples：提供各模块的示例
• models：提供预训练模型，方便快速上手
• tests：单元测试
• tinynn：模型压缩相关代码
  • graph：用于计算图捕获、解析、量化、代码生成、掩码管理等的基础模块
  • prune：剪枝算法
  • converter：模型转换器
  • util：工具类

路线图

• 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 的复杂性和频繁更新，我们无法通过自动化测试覆盖所有情况。 当您遇到问题时，可以查看 FAQ，或通过下方二维码加入钉钉问答群。

TinyNeuralNetwork 快速上手指南

TinyNeuralNetwork 是一个高效且易用的深度学习模型压缩框架，支持神经架构搜索（NAS）、剪枝、量化、模型转换等功能。该框架已广泛应用于天猫精灵、海尔电视、优酷视频及人脸识别考勤机等设备，赋能超过千万台 IoT 设备。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统：Linux / macOS / Windows
• Python 版本：>= 3.9
• 深度学习框架：PyTorch >= 1.10
• 其他依赖：Git

提示：如果您希望使用隔离环境进行开发，也可以直接使用 Docker 构建（需安装 Docker）：

sudo docker build -t tinynn:pytorch1.9.0-cuda11.1 .

安装步骤

您可以选择以下任意一种方式进行安装：

方式一：源码安装（推荐）

克隆仓库并执行安装脚本：

git clone https://github.com/alibaba/TinyNeuralNetwork.git
cd TinyNeuralNetwork
python setup.py install

方式二：pip 一键安装

如果您不想克隆源码，可以直接通过 pip 安装：

pip install git+https://github.com/alibaba/TinyNeuralNetwork.git

国内加速建议：如果下载速度较慢，建议在 pip 命令后添加国内镜像源参数，例如使用阿里云源：

pip install git+https://github.com/alibaba/TinyNeuralNetwork.git -i https://mirrors.aliyun.com/pypi/simple/

基本使用

TinyNeuralNetwork 的核心功能围绕计算图捕获、依赖解析、剪枝和量化展开。以下是一个典型的使用流程示例，展示如何加载预训练模型并进行剪枝准备。

1. 导入模块与加载模型

首先从 tinynn 包中导入必要的工具类，并加载一个预训练的 PyTorch 模型（以 ResNet18 为例）：

import torch
from tinynn.graph.tracer import model_tracer
from tinynn.prune import get_pruner

# 加载预训练模型 (此处以 torchvision 的 resnet18 为例)
model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet18', pretrained=True)
model.eval()

# 创建虚拟输入数据用于追踪计算图
dummy_input = torch.randn(1, 3, 224, 224)

# 使用 Graph Tracer 捕获计算图连接关系
with torch.no_grad():
    graph = model_tracer(model, dummy_input)

2. 配置并执行剪枝

框架内置了多种剪枝算法（如 OneShot L1/L2/FPGM, ADMM, NetAdapt 等）。以下演示如何使用 FPGM 算法进行一次性剪枝：

# 配置剪枝策略
prune_config = {
    'type': 'oneshot',
    'algorithm': 'fpgm',
    'compress_ratio': 0.5,  # 目标压缩率
}

# 获取剪枝器实例
pruner = get_pruner(prune_config, model, graph)

# 执行剪枝
pruner.prune()

# 应用剪枝掩码到模型
pruner.apply_mask()

3. 量化感知训练 (QAT)

TinyNeuralNetwork 自动化了算子融合与计算图量化过程，简化了 QAT 流程：

from tinynn.graph.quantization import quantize

# 对模型进行量化感知训练准备
# 框架会自动处理算子融合和伪量化节点的插入
quantized_model = quantize(model, dummy_input)

# 接下来即可使用标准 PyTorch 流程进行微调训练
# optimizer = ...
# for data, target in loader:
#     ...

4. 模型转换

完成压缩和量化后，可将模型转换为 TFLite 格式以便在端侧部署：

from tinynn.converter import TFLiteConverter

# 初始化转换器
converter = TFLiteConverter(model, dummy_input, tflite_path='model.tflite')

# 执行转换
converter.run()

更多详细用例请参考项目中的 examples 目录。