68 lines
1.8 KiB
Python
68 lines
1.8 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
自定义层量化类型示例
|
|
|
|
场景: 模型整体使用 asymu8 量化, 但指定输入/输出层使用更高精度的量化类型.
|
|
|
|
用法:
|
|
cd examples/custom_layer_quant
|
|
python3 custom_layer_quant.py
|
|
"""
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
from netrans import Netrans
|
|
|
|
MODEL_SRC = os.path.join(os.path.dirname(__file__), "..", "caffe", "lenet_caffe")
|
|
MODEL_SRC = os.path.abspath(MODEL_SRC)
|
|
|
|
|
|
def demo_single_layer():
|
|
"""演示: 单层自定义量化 (正常工作)"""
|
|
print("=" * 60)
|
|
print("场景A: 单个输入层使用 dfpi16, 全模型使用 asymu8")
|
|
print("=" * 60)
|
|
|
|
model = Netrans()
|
|
model.load(MODEL_SRC, mean=[128], std=[1])
|
|
|
|
# input_0 使用 dfpi16, 其余层使用 asymu8
|
|
model.quantize("asymu8", lid="input_0", in_out_quantized="dfpi16")
|
|
model.export("asymu8", platform="pnna")
|
|
|
|
nb = Path(MODEL_SRC) / "wksp" / "lenet_caffe_asymu8_nbg_unify" / "network_binary.nb"
|
|
if nb.exists():
|
|
print(f"OK: {nb} ({nb.stat().st_size:,} bytes)")
|
|
else:
|
|
print("FAIL: NBG 未生成")
|
|
|
|
|
|
def demo_dump_layers():
|
|
"""演示: 查看模型中的层 ID"""
|
|
print("\n" + "=" * 60)
|
|
print("查看模型层信息: netrans dump")
|
|
print("=" * 60)
|
|
|
|
model = Netrans()
|
|
model.load(MODEL_SRC, mean=[128], std=[1])
|
|
model.dump("float32")
|
|
print("dump 完成, 检查 wksp/ 目录获取详细层信息")
|
|
print("关键层: input_0 (输入), output_9 (输出)")
|
|
|
|
|
|
def main():
|
|
if not os.path.isdir(MODEL_SRC):
|
|
print(f"错误: 模型目录不存在 — {MODEL_SRC}")
|
|
sys.exit(1)
|
|
|
|
demo_dump_layers()
|
|
demo_single_layer()
|
|
|
|
print("\n提示: 如需演示多层自定义量化 (lid='input_0,output_9'),")
|
|
print("该功能当前存在已知 bug (BUG-2), 待修复.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|