TPU-MLIR

中文版请见：README。

TPU-MLIR 是一个基于 MLIR 的开源机器学习编译器，专为 TPU 设计。本项目提供了一套完整的工具链，能够将不同框架的预训练神经网络转换为可在 TPU 上高效运行的二进制文件 bmodel。

目前支持的深度学习框架包括 PyTorch、ONNX、TFLite 和 Caffe。其他框架的模型需要先转换为 ONNX 模型。

该项目还支持编译 HuggingFace 的 LLM 模型。目前支持 qwen2 和 llama 系列，未来将增加更多类型的 LLM 模型。

如何构建

首先，你需要安装指定的 Docker，然后可以选择使用预构建包或从源码编译。

安装 Docker

• 从 dockerhub 下载所需镜像。

docker pull sophgo/tpuc_dev:latest

• 如果拉取失败，可以通过以下方式下载：

wget https://sophon-assets.sophon.cn/sophon-prod-s3/drive/25/04/15/16/tpuc_dev_v3.4.tar.gz
docker load -i tpuc_dev_v3.4.tar.gz

• 创建容器：

# myname1234 仅为示例，你可以设置自己的名称
docker run --privileged --name myname1234 -v $PWD:/workspace -it sophgo/tpuc_dev:latest

使用预构建包（方法一）

我们提供了 TPU-MLIR 的 Python 包，可以跳过构建步骤快速安装。
环境要求：python >= 3.10 且 ubuntu:22.04（建议直接使用我们的 Docker 镜像）。

pip install tpu_mlir

编译源码（方法二）

容器创建后，Docker 内的代码目录应为 /workspace/tpu-mlir。

• 构建

在项目目录下运行以下命令：

cd tpu-mlir
source ./envsetup.sh
./build.sh

使用示例（以 Qwen2.5-VL 为例）

以 Qwen2.5-VL 为例，展示如何编译一个 HuggingFace 的 LLM 模型。

1) 下载 Qwen2.5-VL 模型：

git lfs install
git clone git@hf.co:Qwen/Qwen2.5-VL-3B-Instruct-AWQ

2) 在 Docker 环境中编译 Qwen2.5-VL：

llm_convert.py -m /workspace/Qwen2.5-VL-3B-Instruct-AWQ -s 2048 -c bm1684x --max_pixels 672,896 -o qwen2.5vl_3b

llm_convert.py 支持的主要参数如下：

转换完成后，将在指定的输出目录中生成对应的 bmodel 文件。
此示例中的模型已经过量化，因此无需设置 --quantize 参数。

3) 在 PCIe 或 SoC 环境中运行 bmodel：

• 将 python_demo 文件夹复制到你的 PCIe/SoC 环境中，然后构建：

shell mkdir build && cd build cmake .. make cp *cpython*.so .. cd ..

• 将生成的 bmodel 复制到运行时环境并执行：

shell # 将 xxxx.bmodel 替换为实际的 bmodel 文件名 python3 pipeline.py -m xxxx.bmodel -c config

执行结果如下所示：

使用示例（以 yolov5s 为例）

通过一个简单的示例介绍 TPU-MLIR 的用法：编译 yolov5s.onnx 并在 BM1684X TPU 平台上运行。

模型来自 yolov5 官网：https://github.com/ultralytics/yolov5/releases/download/v6.0/yolov5s.onnx。

该模型已放置在项目路径 regression/model/yolov5s.onnx 中。

准备工作

首先，在与本项目同级的目录下创建一个 model_yolov5s 目录，然后将模型和图像文件放入其中。

操作如下：

mkdir model_yolov5s && cd model_yolov5s
cp ${REGRESSION_PATH}/model/yolov5s.onnx .
cp -rf ${REGRESSION_PATH}/dataset/COCO2017 .
cp -rf ${REGRESSION_PATH}/image .
mkdir workspace && cd workspace

模型转换为 MLIR

如果模型以图像作为输入，在转换前需要了解其预处理过程。如果输入是 npz 文件，则无需考虑预处理。预处理公式如下：

$$
y = （x - mean） \times scale,
$$

其中 x 代表输入。

官方 yolov5 的输入是 RGB 图像。每个值都会乘以 1/255。均值和缩放系数分别为 0.0, 0.0, 0.0 和 0.0039216, 0.0039216, 0.0039216。

模型转换命令：

model_transform.py \
    --model_name yolov5s \
    --model_def ../yolov5s.onnx \
    --input_shapes [[1,3,640,640]] \
    --mean 0.0,0.0,0.0 \
    --scale 0.0039216,0.0039216,0.0039216 \
    --keep_aspect_ratio \
    --pixel_format rgb \
    --output_names 350,498,646 \
    --test_input ../image/dog.jpg \
    --test_result yolov5s_top_outputs.npz \
    --mlir yolov5s.mlir

model_transform.py 的主要参数（完整信息请查阅技术参考手册）：

转换为 mlir 文件后，将生成一个包含预处理输入的 ${model_name}_in_f32.npz 文件。

MLIR 转换为 F16 bmodel

通过以下命令将 mlir 文件转换为 F16 bmodel：

model_deploy.py \
  --mlir yolov5s.mlir \
  --quantize F16 \
  --processor bm1684x \
  --test_input yolov5s_in_f32.npz \
  --test_reference yolov5s_top_outputs.npz \
  --model yolov5s_1684x_f16.bmodel

model_deploy.py 的主要参数（完整信息请查阅技术参考手册）：

MLIR 转换为 INT8 bmodel

在转换为 INT8 模型之前，需要运行校准以获得校准表。根据情况，输入数据的数量约为 100 到 1000 张。

然后使用校准表生成对称的 int8 bmodel。如果对称模型已满足要求，通常不建议使用非对称模型，因为非对称模型的性能会略差于对称模型。

以下使用 COCO2017 中的 100 张现有图像进行校准的示例：

run_calibration.py yolov5s.mlir \
  --dataset ../COCO2017 \
  --input_num 100 \
  -o yolov5s_cali_table

执行以下命令转换为 INT8 对称量化模型：

model_deploy.py \
  --mlir yolov5s.mlir \
  --quantize INT8 \
  --calibration_table yolov5s_cali_table \
  --processor bm1684x \
  --test_input yolov5s_in_f32.npz \
  --test_reference yolov5s_top_outputs.npz \
  --tolerance 0.85,0.45 \
  --model yolov5s_1684x_int8.bmodel

结果对比

本项目有一个用 Python 编写的 yolov5 示例（路径：python/samples/detect_yolov5.py），用于目标检测。阅读代码以了解模型的使用方法：
1. 预处理输入
2. 模型推理获取输出
3. 后处理输出

以下代码分别用于验证 onnx/f32/int8 模型的输出：

• ONNX 模型：

detect_yolov5.py \
  --input ../image/dog.jpg \
  --model ../yolov5s.onnx \
  --output dog_origin.jpg

• F16 bmodel：

detect_yolov5.py \
  --input ../image/dog.jpg \
  --model yolov5s_1684x_f16.bmodel \
  --output dog_f16.jpg

• INT8 对称量化 bmodel：

detect_yolov5.py \
  --input ../image/dog.jpg \
  --model yolov5s_1684x_int8.bmodel \
  --output dog_int8.jpg

不同模型的输出对比如下：

辅助工具

模型推理工具 model_runner.py

支持 bmodel/mlir/pytorch/onnx/tflite/caffe。

model_runner.py \
  --input resnet18_in_f32.npz \
  --model resnet18_1684x_f32.bmodel \
  --output resnet18_output.npz

bmodel 工具

bmodel 文件可以通过 model_tool 查看和编辑：

  model_tool
    --info model_file : 显示简要模型信息
    --print model_file : 显示详细模型信息
    --extract model_file : 将多网络 bmodel 提取为多个单网络 bmodel
    --combine file1 .. fileN -o new_file: 通过文件路径合并 bmodel
    --combine_dir dir1 .. dirN -o new_dir: 通过目录路径合并 bmodel
    --dump model_file start_offset byte_size out_file: 从 bmodel 中提取二进制数据到文件

例如，获取 bmodel 的基本信息：

model_tool --info resnet18_1684x_f32.bmodel

资源

以下资源可帮助你更好地理解本项目：