A tool for ONNX model:
- 解析ONNX模型并且编辑: 常量层折叠, Ops fusion.
- 模型分析:Tensor形状推理,每个Op的MACs统计
- Compute Graph 和 Shape Engine.
- 内存压缩:激活Tenosr的内存压缩和权重的内存压缩
- 支持量化模型和稀疏模型.
支持的模型有:
- NLP: BERT, T5, GPT, LLaMa, MPT(TransformerModel)
- Diffusion: Stable Diffusion(TextEncoder, VAE, UNET)
- CV: BEVFormer, MobileNet, YOLO, ...
- Audio: sovits, LPCNet
你可以用onnx_tool.Model类去加载任意ONNX模型,变成易于编辑的python类实例,你可以:
用onnx_tool.Graph类去改变图结构;
用onnx_tool.Node类去改变每个Op的属性和输入输出Tensor;
用onnx_tool.Tensor改变任意Tensor的数据类型和数据内容.
修改完成后,只需要调用Graph或者Model类的save_model接口可以保存所有的修改内容到新的ONNX模型.
每个模型分析报告需要基于某个特定的输入Tensor的形状。所以在分析模型之前要先进行一次形状推理。
浮点乘加数(等于2倍的浮点操作数), 内存占用(字节数), 参数量(参数个数)稀疏的块的形状, 稀疏块的稀疏率(全为0的稀疏块的稀疏率), 参数的稀疏率(数值为0的稀疏率)
how to use: data/Profile.md.
pytorch usage: data/PytorchUsage.md.
tensorflow
usage: data/TensorflowUsage.md.
examples: benchmark/examples.py.
移除了所有的Tensor形状计算op, 更新动态Tensor的形状可以用Shape Engine来替代。推理引擎只需要负责计算图的计算,不需要考虑Tensor的形状更新。
examples:
benchmark/shape_regress.py.
benchmark/examples.py.
如何集成 Compute Graph 和 Shape Engine 到cpp推理引擎中: data/inference_engine.md
MHA and Layernorm Fusion for Transformers
Resnet18 fusionhow to use: data/Subgraph.md.
BERT examples: benchmark/examples.py.
Pattern fusion: benchmark/do_fusion.py.
可以帮助实现model parallel。
how to use: data/Subgraph.md.
对于LLM和高分辨CV模型, 激活内存的压缩可以帮助节省整个模型的内存使用.
压缩方法可以在大多数模型上实现 5% 内存压缩率.
例如:
model | Native Memory Size(MB) | Compressed Memory Size(MB) | Compression Ratio(%) |
---|---|---|---|
StableDiffusion(VAE_encoder) | 14,245 | 540 | 3.7 |
StableDiffusion(VAE_decoder) | 25,417 | 1,140 | 4.48 |
StableDiffusion(Text_encoder) | 215 | 5 | 2.5 |
StableDiffusion(UNet) | 36,135 | 2,232 | 6.2 |
GPT2 | 40 | 2 | 6.9 |
BERT | 2,170 | 27 | 1.25 |
code example: benchmark/compression.py
pip install onnx-tool
OR
pip install --upgrade git+https://github.com/ThanatosShinji/onnx-tool.git
python>=3.6
If pip install onnx-tool
failed by onnx's installation, you may try pip install onnx==1.8.1
(a lower version like this) first.
Then pip install onnx-tool
again.
- Loop op is not supported
Results of ONNX Model Zoo and SOTA models
注意对于支持动态输入形状的模型,模型的MACs随输入形状的改变而改变。下表中的MACs数据是基于data/public/config.py中的配置输入形状得到。 带有所有Tensor形状的模型和分析报告可以从下面的网盘中下载: baidu drive(code: p91k) google drive
|
|