如何蒸馏模型并部署到ARM 开发板 之二 代码实现

软硬件操作环境准备

将大模型蒸馏并部署到高通骁龙ARM板卡上是一个系统工程，涉及模型压缩、转换和边缘计算。训练和部署项目代码在文章最后。

软硬件操作环境准备

在开始之前，需要准备好以下环境。这是项目成功的基础。

详细执行步骤

整个过程分为五个主要阶段：准备、蒸馏、转换、部署、测试。

第一阶段：环境与数据准备

第一步：搭建开发环境

1. 在开发机上安装Ubuntu或配置WSL2（在Windows11 系统上运行Ubuntu）。
2. 安装conda(anaconda安装教程自行查询)，并创建一个新的Python环境：conda create -n model_distillation python=3.12
3. 激活环境：conda activate model_distillation
4. 安装PyTorch（参考官网命令，选择适合您CUDA版本的安装命令）。、
5. 安装Hugging Face库：pip install transformers datasets
6. 安装蒸馏库，例如TextBrewer：pip install textbrewer
7. 下载并解压高通SNPE SDK。按照官方文档安装Python依赖并配置环境变量（如SNPE_ROOT）
8. 安装辅助库，用于数据处理和评估 pip install tqdm scikit-learn pandas
9. 创建项目目录：mkdir model_distillation && cd model_distillation

第二步：获取教师模型和学生模型

脚本里面自动下载。

• 教师模型： 选择一个大型模型，如bert-base-chinese（用于中文）或bert-base-uncased（用于英文）。
  • 网络考虑： 国内从Hugging Face官网下载可能较慢。解决方案：
    • 使用镜像站： 在代码中设置HF_ENDPOINT=https://hf-mirror.com环境变量，或使用huggingface-cli download --resume-download --local-dir-use-symlinks False命令配合镜像站下载。
    • 手动下载： 从Hugging Face Hub网页下载所有文件，然后用from_pretrained(‘/local/path’)加载。
• 学生模型： 选择一个结构更小、层数更少的模型架构。例如，目标是蒸馏出一个6层的BERT，学生模型可以初始化为一个6层的BERT结构（如huggingface提供的bert-base-uncased配置，但修改层数为6）。

第三步：准备蒸馏数据
蒸馏需要一份训练数据。可以使用：

• 公开数据集（如GLUE、SQuAD等，可通过datasets库加载）。
• 业务领域的特定无标签数据。对于知识蒸馏，无标签数据就足够了，教师模型会为其生成“软标签”（概率分布）。

将使用Hugging Face datasets 库直接下载IMDb数据集，无需手动下载。如果网络不畅，可以设置镜像。数据下载会在后续脚本中自动完成。

# 在终端中设置HF镜像（可选，国内用户建议使用）
export HF_ENDPOINT=https://hf-mirror.com

第二阶段：知识蒸馏

第四步：实现蒸馏流程
以下是使用TextBrewer的伪代码逻辑流程：

1. 加载模型：

   # 训练脚本 distill_train.py
   from transformers import AutoModelForSequenceClassification, AutoTokenizer
   from textbrewer import GeneralDistiller, TrainingConfig, DistillationConfig
   
   teacher_model = AutoModelForSequenceClassification.from_pretrained('bert-base-uncased')
   student_model = AutoModelForSequenceClassification.from_pretrained('./student-initial-config') # 你的小模型
   
   # 将教师模型设置为评估模式，不更新其参数
   teacher_model.eval()
   

2. 定义蒸馏配置：

   # 训练配置 config.py
   train_config = TrainingConfig(
       output_dir = './distillation_output',
       ckpt_frequency = 1000, # 保存检查点频率
       log_dir = './log',
   )
   
   # 蒸馏配置（核心）
   distill_config = DistillationConfig(
       temperature = 8, # 温度参数，软化教师输出，蕴含更多知识
       hard_label_weight = 0.1, # 真实标签的权重（如果可用）
       soft_label_weight = 0.9, # 教师软标签的权重
       kd_loss_type = 'ce', # 知识蒸馏损失函数类型（交叉熵）
   )
   

3. 定义适应函数： 这是一个关键函数，用于告诉蒸馏器如何从模型的输出中提取需要的logits（预测分数）。

   # 训练脚本 distill_train.py
   def simple_adaptor(batch, model_outputs):
       # 假设你的模型输出是一个包含logits的元组或类
       return {'logits': model_outputs.logits}
   

4. 创建蒸馏器并开始训练：

   # 训练脚本 distill_train.py
   distiller = GeneralDistiller(
       train_config = train_config,
       distill_config = distill_config,
       model_T = teacher_model, # Teacher
       model_S = student_model, # Student
       adaptor_T = simple_adaptor,
       adaptor_S = simple_adaptor
   )
   
   # 假设你有一个DataLoader `train_dataloader`
   distiller.train(optimizer, train_dataloader, num_epochs=3)
   

   # 确保在 ‘model_distillation’ 目录下
   python distill_script.py
   

第三阶段：模型转换

第五步：将蒸馏后的模型转换为ONNX
在ARM开发板的Ubuntu系统上使用蒸馏后的模型，一个很好的思路是将模型转换为更轻量级、更适合嵌入式设备的格式。这里提供三种主要的部署方式，可以根据需求选择：

方案一：使用ONNX Runtime（推荐首选）

这个方案兼容性好，社区支持完善，是大多数情况下的首选。

1. 转换模型为ONNX格式
   在你的开发电脑（拥有完整PyTorch环境）上，运行以下脚本，将蒸馏出的模型转换为.onnx文件。

   # convert2onnx.py
   import torch
   from transformers import AutoModelForSequenceClassification, AutoTokenizer
   
   # 加载你训练好的模型和分词器
   model_path = "./models/distilled_student"
   model = AutoModelForSequenceClassification.from_pretrained(model_path)
   tokenizer = AutoTokenizer.from_pretrained(model_path)
   
   # 设置为评估模式
   model.eval()
   
   # 创建一个示例输入（dummy input）
   # 注意：这里的形状需要与你模型训练时一致
   dummy_input = torch.randint(0, tokenizer.vocab_size, (1, 256))  # 示例: (batch_size, sequence_length)
   
   # 指定输入和输出的名称
   input_names = ["input_ids"]
   output_names = ["logits"]
   
   # 导出模型
   torch.onnx.export(model,
                     dummy_input,
                     "distilled_model.onnx",
                     input_names=input_names,
                     output_names=output_names,
                     dynamic_axes={'input_ids': {0: 'batch_size'}, 'logits': {0: 'batch_size'}},  # 支持动态batch
                     opset_version=14,  # 使用较新的ONNX算子集
                     export_params=True)
   print("模型已成功导出为 distilled_model.onnx")
   

   # 运行转换脚本
   python convert2onnx.py
   

   将生成的 distilled_model.onnx 文件拷贝到你的ARM开发板中。

方案二：使用超轻量级推理框架TinyMaix

如果你的开发板资源非常有限，或者模型非常简单，可以尝试这个方案。

1. 安装依赖与编译TinyMaix
   在开发板上操作：

   # 安装基础编译工具
   sudo apt update
   sudo apt install cmake gcc make git
   # 克隆TinyMaix仓库
   git clone https://github.com/sipeed/TinyMaix.git
   cd TinyMaix
   # 编译（假设是ARM64 Linux平台）
   mkdir build && cd build
   cmake ..
   make
   

   编译成功后，你可以在 examples 目录下找到一些示例程序。

2. 模型转换与使用
   TinyMaix支持从Keras H5或TFLite模型转换。你需要先将你的PyTorch模型转换为这两种格式之一，然后再使用TinyMaix提供的转换工具生成其支持的模型格式。由于转换步骤相对复杂，且涉及多个框架，建议你参考 TinyMaix官方GitHub页面 获取最新的转换指南。

方案三：使用针对Arm优化的Tengine

如果你的ARM开发板带有Mali GPU或者NPU，这个方案可能带来更好的性能。

1. 编译Tengine
   在开发板上操作：

   # 安装依赖
   sudo apt install libprotobuf-dev protobuf-compiler libopencv-dev pkg-config
   # 克隆Tengine仓库，注意参数
   git clone --recurse-submodules https://github.com/OAID/Tengine.git
   cd Tengine
   # 使用提供的配置脚本进行编译
   ./linux_build.sh default_config/arm64_linux_native.config
   

   编译完成后，推理相关的库文件和示例会在 build 目录下生成。

2. 模型转换与部署
   Tengine支持将多种格式的模型转换为其专属格式。你需要使用其提供的 convert_tool 工具将你的模型（例如ONNX格式）进行转换。之后，你可以参考Tengine提供的C++或C API示例来编写你的推理程序。

第四阶段：部署到骁龙板卡

第六步：准备部署环境

在ARM开发板上安装ONNX Runtime

在开发板的Ubuntu系统上，需要安装针对ARM64架构的ONNX Runtime。访问 ONNX Runtime官方GitHub 的发布页面，查找适用于Linux ARM64的版本。
例如，使用pip安装：

pip install onnxruntime

第七步：编写C++或Python推理应用程序

在开发板上创建新的推理脚本，例如 onnx_inference.py。

# inference_onnx.py
import onnxruntime as ort
import numpy as np
from transformers import AutoTokenizer

# 加载分词器
tokenizer = AutoTokenizer.from_pretrained("./models/distilled_student")

# 创建ONNX Runtime推理会话
session = ort.InferenceSession("distilled_model.onnx")

def predict(text):
    # 预处理文本
    inputs = tokenizer(text, return_tensors="np", truncation=True, padding=True, max_length=256)
    input_ids = inputs["input_ids"].astype(np.int64)
    
    # 运行模型推理
    outputs = session.run(["logits"], {"input_ids": input_ids})
    logits = outputs[0]
    
    # 后处理：获取预测结果
    predicted_class_id = np.argmax(logits, axis=1)[0]
    confidence = np.max(logits, axis=1)[0]
    
    return {"predicted_class": predicted_class_id, "confidence": confidence}

# 使用示例
if __name__ == "__main__":
    test_text = "This movie is absolutely fantastic!"
    result = predict(test_text)
    print(f"文本: '{test_text}'")
    print(f"预测类别: {result['predicted_class']}")
    print(f"置信度: {result['confidence']:.4f}")

第八步：交叉编译和运行

1. 后续将模型封装成API或者库文件，供后续调用。

# 运行推理脚本，看日志输出模型结果
python inference_onnx.py

第五阶段：测试与优化

第九步：功能与性能测试

1. 功能测试： 输入一些测试文本，确保模型的输出与在开发机上的Python环境中的结果基本一致（允许有细微的数值误差）。
2. 性能分析： 使用SNPE提供的分析工具（如snpe-diagview）来评估模型在板卡上的推理速度和内存占用。

第十步：迭代优化
如果性能不达标，可以考虑：

• 量化： 使用SNPE的量化工具将FP32模型转换为INT8模型，能显著减小模型体积并提升速度。
• 尝试其他运行时： 将模型部署到骁龙的Hexagon DSP或Adreno GPU上，通常能获得比CPU更好的性能。但这可能需要额外的步骤（如量化、DSP签名等）。

代码

目录结构

1. 开发机

model_distillation
|-config.py
|-convert2onnx.py
|-distill_train.py
|-utils.py
|-models/distilled_student
   |-config.json
   |-pytorch_model.bin
   |-special_tokens_map.json
   |-tokenizer_config.json
   |-tokenizer.json
   |-vocab.txt
   |-distilled_model_dynamic.onnx

目录结构说明：

蒸馏后的模型： 保存在 ./models/distilled_student 目录下，包含：

• config.json (模型架构配置文件)：定义模型的层数、隐藏层大小、注意力头数等结构信息
• pytorch_model.bin (学生模型权重) ：包含所有训练好的参数和权重
• special_tokens_map.json(特殊token映射)：定义[CLS]、[SEP]、[PAD]等特殊token
• tokenizer.json (分词器文件)：定义如何将文本转换为模型可理解的token
• vocab.txt (词汇表)：包含所有token及其对应的ID
• distilled_model_dynamic.onnx(转化后的模型，可以部署到arm板子)

distilled_student整个目录下的内容，就是可以直接使用的模型，加载这个目录即调用模型，最后有一个onnx文件，是因为需要在arm板子上使用，转化后才生成的文件。

2. 部署机

distilled_student
|-inference_onnx.py
|-distilled_Student
   |-distilled_model_dynamic.onnx

目录结果说明：

一个模型，一个使用模型的例子

• inference_onnx.py 使用模型的列子，演示如何使用模型
• distilled_model_dynamic.onnx 模型文件

代码

# 记得star 哦 ❤️👍😘
git@github.com:NoelCarlton/distill-fisrt.git

总结

这个过程技术链条较长，每一步都可能遇到挑战。建议从一个非常小的模型（如蒸馏一个两层的BERT在简单任务上）开始，走通整个流程，然后再应用到您实际的大模型上。重点关注模型转换（ONNX） 和板卡端环境配置这两个最容易出错的环节。