模型批量时间测试工具

概述

本工具通过 Python 接口实现板卡推理，批量测试每个模型的多种量化类型板卡推理时间，结果保存到 CSV 文件中。支持多输入模型的自动化测试。

流程图

graph TD
    subgraph 模型批量处理流程
        A[开始] --> B[板卡建立WebSocket连接]
        B --> C[清空CSV文件
写入表头]
        C --> D[遍历模型文件夹
获取所有模型]
        D --> E{遍历每个模型}

        subgraph E[单个模型处理流程]
            E1[查找inputmeta配置文件] --> E2[解析YAML
生成随机输入数据]
            E2 --> E3[获取量化文件夹列表]
            E3 --> E4{遍历每个量化类型}
        end

        subgraph F[单次推理测试流程]
            F1[删除旧模型] --> F2[加载新量化模型]
            F2 --> F3[执行多次推理]
            F3 --> F4[统计推理时间
计算平均和最佳时间]
            F4 --> F5[追加结果到CSV]
        end

        E4 --> F
        F --> G{是否还有
下一个量化类型?}
        G -- 是 --> E4
        G -- 否 --> H{是否还有
下一个模型?}
        H -- 是 --> E
        H -- 否 --> I[关闭连接
结束]
    end

项目结构

├── resource/             # 算法模型
├── src/                  # 源码
│   ├── utils/            # 公共组件
│   │   └── client.py     # PNNA设备客户端
│   └── auto_test.py      # 批量测试脚本
├── README.md             # 项目说明文档
└── requirements.txt      # 依赖包

快速开始

环境要求及配置

硬件要求

NPU 板卡：FT78E_DSKC_202311V0 板卡
连接与通信准备（自备）：网线
板卡上已开启 jsonrpc_server 服务

编译和使用

下载下位机源码：下位机源码路径工程下的 examples/jsonrpc_server 下
编译工具和相关库请看 jsonrpc_server 的页面
将板子 IP 设置为 192.168.23.100，在电脑上添加 192.168.23.132 这个 IP（与板子 IP 保持在同一网段）
将可执行文件（examples/jsonrpc_server 下的 server）传输到板子，在命令行输入 ./server 运行即可

通信基础信息

通信协议：WebSocket
连接地址：ws://192.168.23.100:8000/websocket，与板卡保持一致
数据格式：JSON-RPC（远程过程调用协议，请求/响应均为 JSON 格式）

创建 Python 3.8 环境

wget "https://mirrors.tuna.tsinghua.edu.cn/github-release/conda-forge/miniforge/LatestRelease//Miniforge3-$(uname)-$(uname -m).sh"
mkdir -p ~/app
INSTALL_PATH="${HOME}/app/miniforge3"
bash Miniforge3-Linux-x86_64.sh -b -p ${INSTALL_PATH}
echo "source "${INSTALL_PATH}/etc/profile.d/conda.sh"" >> ${HOME}/.bashrc
echo "source "${INSTALL_PATH}/etc/profile.d/mamba.sh"" >> ${HOME}/.bashrc
source ${HOME}/.bashrc
mamba create -n netrans python=3.8 -y
mamba activate netrans

安装依赖

pip install -r requirements.txt

下载工程

git clone https://gitlink.org.cn/nudt_dsp/models.git

使用方法

执行脚本

python auto_test.py --models_dir models_test

命令行参数

参数	类型	默认值	说明
`--models_dir`	str	`models_test`	模型目录路径
`--output_csv`	str	`results.csv`	输出 CSV 文件路径
`--runs`	int	`3`	每个模型的推理次数
`--host`	str	`192.168.23.100:8000`	PNNA 设备 WebSocket 地址

输出结果

生成的 CSV 文件包含以下列：

字段名	说明
model_name	模型名称
model_type	量化类型
avg(ms)	平均推理时间（毫秒）
best(ms)	最佳推理时间（毫秒）

评估指标

推理平均速度：下位机多次推理平均时间（毫秒）
推理最佳速度：下位机单次推理最佳时间（毫秒）

函数说明

clear_csv_file(csv_file)

清空 CSV 文件，只保留表头。

参数：

csv_file (str): CSV 文件路径

append_to_csv(result, csv_file)

追加单条结果到 CSV 文件。

参数：

result (dict): 包含测试结果的数据字典
csv_file (str): CSV 文件路径

返回：

bool: 写入是否成功

process_nb_file(pnna_device, nb_path, model_name, quant_type, input_data_list, num_runs)

处理单个 .nb 文件并返回测试结果。

参数：

pnna_device (Client): PNNA 设备客户端
nb_path (str): .nb 模型文件路径
model_name (str): 模型名称
quant_type (str): 量化类型
input_data_list (list): 输入数据列表，每个元素是一个 numpy 数组
num_runs (int): 推理次数

返回：

dict 或 None: 测试结果字典，错误时返回 None

generate_data_from_yml(yml_path)

从 YAML 配置文件中读取所有输入数据的 shape 并生成 uint8 随机数据。

参数：

yml_path (str): YAML 配置文件路径

返回：

list 或 None: 输入数据列表，每个元素是一个 numpy 数组

find_inputmeta_yml(model_path)

在模型目录中查找 inputmeta 配置文件。

参数：

model_path (str): 模型目录路径

返回：

str 或 None: 配置文件路径，未找到返回 None

get_quant_folders(wksp_path)

获取量化文件夹列表并按指定顺序排序。

排序优先级：asymu8 → symi8 → dfpi16 → fp16

参数：

wksp_path (str): wksp 目录路径

返回：

list: 排序后的量化文件夹名称列表

process_and_save_model(pnna_device, model_path, model_name, csv_file, num_runs)

处理单个模型的所有量化文件，并实时保存到 CSV。

参数：

pnna_device (Client): PNNA 设备客户端
model_path (str): 模型目录路径
model_name (str): 模型名称
csv_file (str): 输出 CSV 文件路径
num_runs (int): 推理次数

返回：

int: 成功处理的配置数量

process_models_incremental(pnna_device, models_dir, output_csv, num_runs)

增量处理所有模型，处理完一个模型就保存一次。

参数：

pnna_device (Client): PNNA 设备客户端
models_dir (str): 模型根目录
output_csv (str): 输出 CSV 文件路径
num_runs (int): 推理次数

main()

主函数：解析命令行参数并启动测试。

作者: liangliangou