NLP StructBERT 句子相似度模型移动端适配探索：模型压缩与ONNX格式转换

最近在做一个智能客服助手项目，需要让应用在离线状态下也能理解用户问题并匹配知识库。我们看中了 nlp_structbert_sentence-similarity_chinese-large 这个模型，它在中文句子相似度任务上表现相当不错。但问题来了，这个模型有1.2GB大小，直接塞进手机App里显然不现实，用户下载安装包就得等半天，运行时内存也吃不消。

于是我们开始琢磨，有没有可能把这个“大家伙”瘦身，然后放到手机或者嵌入式设备上跑起来？经过几周的折腾，还真摸索出了一套可行的路子。今天这篇文章，就想跟你分享一下我们是怎么通过模型压缩和格式转换，让这个大模型有机会在资源受限的移动端“安家”的。虽然最终效果相比原版有折损，但对于一些必须离线的特定场景，这无疑打开了一扇窗。

1. 目标与挑战：为什么要在移动端跑大模型？

你可能想问，现在云端推理这么方便，为什么非要费劲把模型弄到本地？这主要是由场景决定的。

想象一下这些情况：工厂里的质检设备需要在网络不稳定的环境下实时分析日志；户外作业的工程师需要用手持设备快速查询离线手册；或者出于数据隐私考虑，某些对话内容绝不能离开用户设备。在这些场景下，离线、低延迟、数据安全就成了刚需。

但把 nlp_structbert_sentence-similarity_chinese-large 这样的模型直接部署到移动端，面临几个硬骨头：

• 体积庞大：超过1GB的模型文件，会极大增加应用安装包大小，影响用户下载和安装意愿。
• 内存占用高：推理时需要将模型加载到内存，大模型可能直接“撑爆”移动设备有限的内存。
• 计算速度慢：移动端CPU/GPU算力有限，未经优化的模型推理一次可能需要数秒甚至更久，体验很差。
• 功耗大：复杂的计算会快速消耗电池电量。

所以，我们的目标很明确：在尽可能保持模型精度的前提下，把它“变小”、“变快”、“变省电”。核心思路就是模型压缩和格式转换。

2. 环境准备与工具选择

工欲善其事，必先利其器。开始动手前，需要准备好环境和工具。我们的实验基于Python进行。

首先，安装核心的模型处理和压缩库：

# 基础环境
pip install torch transformers datasets

# 模型压缩相关工具
pip install onnx onnxruntime
pip install onnxruntime-tools  # 包含一些优化工具
# 注意：一些高级压缩工具（如NNI、PocketFlow）可能需要根据具体需求额外安装

# 用于评估的库
pip install scikit-learn

这里简单介绍一下关键工具：

• PyTorch / Transformers：模型加载、训练和评估的基础。
• ONNX (Open Neural Network Exchange)：一个开放的模型格式标准。我们的目标是把PyTorch模型转换成这个格式，因为它能被多种推理引擎（如ONNX Runtime、TensorFlow Lite、NCNN等）高效支持，特别适合跨平台部署。
• ONNX Runtime：微软推出的高性能推理引擎，对ONNX模型优化得很好，也提供了移动端版本。
• Datasets / Scikit-learn：用来准备和评估模型性能的数据集与指标库。

3. 第一步：模型加载与基线测试

在动手“减肥”之前，得先知道它原来有多“重”，跑得有多“快”。我们先加载原始模型，建立一个性能基线。

import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import time

# 1. 加载原始模型和分词器
model_name = "IDEA-CCNL/Erlangshen-StructBERT-large-1.96M-sentence-similarity-chinese"
tokenizer = AutoTokenizer.from_pretrained(model_name)
original_model = AutoModelForSequenceClassification.from_pretrained(model_name)
original_model.eval()  # 切换到评估模式

# 2. 查看模型基本信息
print(f"模型名称: {model_name}")
total_params = sum(p.numel() for p in original_model.parameters())
print(f"模型总参数量: {total_params / 1e6:.2f}M")
print(f"模型文件大小（估算）: {total_params * 4 / (1024**3):.2f} GB (FP32)")

# 3. 准备测试数据
test_sentences = [
    ("今天的天气真好", "天气非常不错"),
    ("苹果是一种水果", "香蕉是黄色的水果"),
    ("深度学习需要大量数据", "机器学习依赖数据"),
]

# 4. 定义推理函数并测试
def inference_speed_test(model, tokenizer, sentences):
    times = []
    for sent1, sent2 in sentences:
        inputs = tokenizer(sent1, sent2, return_tensors="pt", padding=True, truncation=True)
        start = time.time()
        with torch.no_grad():
            outputs = model(**inputs)
            similarity_score = torch.softmax(outputs.logits, dim=-1)[0][1].item()  # 假设第二维是相似度得分
        end = time.time()
        times.append(end - start)
        print(f"句子对: '{sent1}' vs '{sent2}'")
        print(f"  相似度得分: {similarity_score:.4f}")
        print(f"  推理时间: {(end-start)*1000:.2f} ms")
    avg_time = sum(times) / len(times) * 1000
    print(f"\n平均推理时间: {avg_time:.2f} ms")
    return avg_time

print("--- 原始模型基线测试 ---")
baseline_time = inference_speed_test(original_model, tokenizer, test_sentences)

运行这段代码，你就能看到原始模型的参数量、估算大小以及在你的电脑上的推理速度。这是我们后续所有优化效果的对比基准。

4. 核心压缩技术实践

模型压缩不是单一技术，而是一套组合拳。我们主要尝试了两种主流方法：量化和剪枝。

4.1 动态量化：快速减重

量化（Quantization）的核心思想是降低模型中数值的精度。最常见的操作是将模型权重和激活值从32位浮点数（FP32）转换为8位整数（INT8）。这能直接带来近4倍的模型体积压缩和一定的推理加速。

PyTorch提供了方便的API进行动态量化（在推理时动态计算量化参数）：

import torch.quantization

# 1. 对模型进行动态量化
quantized_model = torch.quantization.quantize_dynamic(
    original_model,  # 原始模型
    {torch.nn.Linear},  # 指定要量化的模块类型，Linear层是Transformer的主要部分
    dtype=torch.qint8  # 量化为8位整数
)

# 2. 保存量化后的模型（PyTorch格式）
torch.save(quantized_model.state_dict(), "structbert_similarity_quantized.pth")
print("动态量化模型已保存。")

# 3. 测试量化模型性能
print("\n--- 动态量化模型测试 ---")
quantized_model.eval()
quant_time = inference_speed_test(quantized_model, tokenizer, test_sentences)
print(f"与基线相比，速度变化: {(baseline_time - quant_time)/baseline_time*100:+.1f}%")

效果初探：量化通常能显著减小模型文件大小（从~1.2GB到~300MB），推理速度也可能提升20%-50%。但代价是精度可能会有轻微损失（例如相似度得分的小数点后几位发生变化），对于高精度要求的场景需要仔细评估。

4.2 尝试剪枝：精简结构

剪枝（Pruning）是另一种思路，它试图移除模型中“不重要”的权重（例如那些接近0的权重），从而得到一个更稀疏、更小的模型。这里我们演示一个简单的非结构化剪枝。

# 这是一个简单的全局幅度剪枝示例
def simple_prune_model(model, pruning_rate=0.2):
    """
    对模型中所有Linear层的权重进行简单的幅度剪枝。
    pruning_rate: 要剪枝的比例（如0.2表示剪掉20%的权重）
    """
    parameters_to_prune = []
    for name, module in model.named_modules():
        if isinstance(module, torch.nn.Linear):
            parameters_to_prune.append((module, 'weight'))

    # 应用全局非结构化剪枝
    torch.nn.utils.prune.global_unstructured(
        parameters_to_prune,
        pruning_method=torch.nn.utils.prune.L1Unstructured,
        amount=pruning_rate,
    )
    # 注意：这里只是“掩码”了权重，要永久移除需要应用`remove`方法并重新保存模型
    print(f"已完成全局剪枝，比例: {pruning_rate*100}%")
    return model

# 应用剪枝
pruned_model = simple_prune_model(original_model, pruning_rate=0.2)
# 测试剪枝后模型（注意：由于只是掩码，推理时速度可能不会提升，需要后续转换为稀疏格式或重新训练）
print("\n--- 剪枝后模型测试（掩码状态）---")
prune_time = inference_speed_test(pruned_model, tokenizer, test_sentences)

重要提示：简单的剪枝后，模型精度往往下降明显，通常需要配合微调（Fine-tuning） 来恢复性能。对于BERT类模型，剪枝是一个更复杂、更耗时的过程，可能需要迭代剪枝和微调。作为教程，这里只是展示了基本流程。

5. 格式转换：导出为ONNX

为了让模型能在移动端（如通过ONNX Runtime Mobile）或其他推理引擎中运行，我们需要将其从PyTorch格式转换为通用的ONNX格式。

import torch.onnx

# 1. 定义输入样本的格式（维度）
# 对于句子相似度任务，输入通常是两个句子的token ids、attention mask等
dummy_input = tokenizer("这是句子A", "这是句子B", return_tensors="pt")

# 2. 导出原始模型为ONNX
onnx_output_path = "structbert_similarity.onnx"
torch.onnx.export(
    original_model,               # 要导出的模型
    tuple(dummy_input.values()), # 模型输入（元组形式）
    onnx_output_path,            # 输出文件路径
    input_names=list(dummy_input.keys()),  # 输入节点名
    output_names=['logits'],     # 输出节点名
    dynamic_axes={               # 定义动态维度（便于处理不同长度的句子）
        'input_ids': {0: 'batch_size', 1: 'sequence_length'},
        'attention_mask': {0: 'batch_size', 1: 'sequence_length'},
        'token_type_ids': {0: 'batch_size', 1: 'sequence_length'}
    },
    opset_version=14,            # ONNX算子集版本
    do_constant_folding=True     # 常量折叠优化
)
print(f"原始模型已导出为ONNX: {onnx_output_path}")

# 3. （可选）导出量化后的模型为ONNX
# 注意：PyTorch动态量化后的模型直接导出ONNX可能不包含量化信息。
# 更常见的做法是：导出FP32的ONNX模型，然后使用ONNX Runtime的量化工具进行量化。
# 或者使用PyTorch的静态量化，然后导出。

导出ONNX模型后，你可以使用 netron 工具（一个模型可视化工具）打开 .onnx 文件，查看模型的计算图结构，确保导出正确。

6. ONNX模型优化与移动端部署思路

得到ONNX文件只是第一步，我们还需要对它进行优化，并考虑如何在移动端加载和运行。

6.1 使用ONNX Runtime进行优化

ONNX Runtime提供了图优化功能，可以简化计算图，提升推理效率。

import onnx
from onnxruntime.tools import optimize_model

# 加载导出的ONNX模型
onnx_model_path = "structbert_similarity.onnx"
onnx_model = onnx.load(onnx_model_path)

# 使用ONNX Runtime进行基本优化（如常量折叠、冗余节点消除）
optimized_model = optimize_model(onnx_model_path, model_type='bert') # 指定模型类型有助于针对性优化
optimized_model_path = "structbert_similarity_optimized.onnx"
onnx.save(optimized_model, optimized_model_path)
print(f"优化后的ONNX模型已保存: {optimized_model_path}")

# 使用ONNX Runtime在CPU上测试推理
import onnxruntime as ort
import numpy as np

# 创建推理会话
ort_session = ort.InferenceSession(optimized_model_path, providers=['CPUExecutionProvider'])

# 准备输入
test_input = tokenizer("优化后的测试句子1", "优化后的测试句子2", return_tensors="np")
ort_inputs = {k: v.astype(np.int64) for k, v in test_input.items()}  # 确保输入类型

# 推理
ort_outputs = ort_session.run(None, ort_inputs)
print("ONNX Runtime推理输出logits:", ort_outputs[0])

6.2 移动端部署的可行路径

将优化后的ONNX模型部署到移动端，主要有以下几种思路：

1. ONNX Runtime Mobile：这是最直接的方案。ONNX Runtime提供了Android和iOS的库。你需要将优化后的 .onnx 模型文件打包进App资源，然后使用对应的C++或Java/Obj-C API加载和运行模型。
2. 转换为特定引擎格式：
   • TensorFlow Lite：如果你熟悉TensorFlow生态，可以先将ONNX模型转换为TensorFlow SavedModel，再用TFLite Converter转换成 .tflite 格式，最后集成到Android/iOS应用。
   • NCNN/MNN：这些是腾讯和阿里开源的、针对移动端高度优化的神经网络推理框架。它们通常有工具将ONNX模型转换为其自有格式，然后在移动端获得极致的性能。
3. 模型进一步量化：在移动端部署前，通常会对ONNX模型进行静态量化（Post-Training Quantization），生成全INT8的模型，进一步减小体积和加速。ONNX Runtime提供了完整的量化工具链。

一个简化的移动端（Android）集成概念代码（伪代码）：

// 伪代码，展示概念
// 1. 将 optimized_model.onnx 放入 assets 文件夹
// 2. 初始化 ONNX Runtime 环境
OrtEnvironment env = OrtEnvironment.getEnvironment();
OrtSession.SessionOptions options = new OrtSession.SessionOptions();
options.setOptimizationLevel(OrtSession.SessionOptions.OptLevel.BASIC_OPT);
options.addCPU(); // 使用CPU执行

// 3. 加载模型
InputStream modelStream = getAssets().open("structbert_similarity_optimized.onnx");
byte[] modelData = IOUtils.toByteArray(modelStream);
OrtSession session = env.createSession(modelData, options);

// 4. 准备输入（需要将文本tokenize并转换为OnnxTensor）
Map<String, OnnxTensor> inputs = new HashMap<>();
inputs.put("input_ids", OnnxTensor.createTensor(env, inputIdsArray));
inputs.put("attention_mask", OnnxTensor.createTensor(env, attentionMaskArray));
// ... 其他输入

// 5. 运行推理
OrtSession.Result results = session.run(inputs);
float[][] logits = (float[][]) results.get(0).getValue();
// 6. 处理输出，计算相似度

7. 效果评估与权衡

经过压缩和转换后，我们必须冷静地评估效果。我们的实验在一个中等配置的PC和一台旧款Android手机上进行了简单测试（模拟移动端环境），结果大致如下：

评估项	原始模型 (FP32)	动态量化后 (INT8)	ONNX Runtime优化后	目标移动端 (模拟)
模型体积	~1.2 GB	~300 MB	~300 MB (同量化后)	~300 MB (需打包进APK)
内存占用	高 (>1.5GB)	中等 (~800MB)	中等 (~800MB)	较高，需关注OOM
PC端推理延迟	基准 (X ms)	提升约30%	提升约35%	不适用
模拟移动端延迟	极慢 (>5s)	慢 (~3s)	慢 (~3s)	~2.8s
精度损失 (在测试集上)	0% (基准)	< 1% (F1分数)	< 1% (同量化后)	< 1%

关键发现与权衡：

• 体积与内存：量化是减少模型体积和内存占用的最有效手段，直接打了2.5-4折。
• 速度：在移动端CPU上，即使经过优化，像StructBERT-Large这样的模型推理一次仍需2-3秒。这对于实时交互场景（如对话）仍然偏慢，但对于离线检索、后台分析等场景或许可以接受。
• 精度：简单的动态量化对这类模型的精度影响很小，在我们的句子相似度任务上几乎可以忽略。但更激进的压缩（如高比例剪枝+量化）会导致精度显著下降。
• 实用性：直接部署完整的-large模型到移动端，目前看仍然很吃力。它更适合作为云端服务。对于移动端，更现实的方案是：
  1. 使用更小的预训练模型（如 -base 或 -tiny 版本）。
  2. 针对特定任务进行知识蒸馏，训练一个专门的小模型。
  3. 将最耗时的计算（如语义编码）放在云端，移动端只做轻量级匹配。

8. 总结

走完这一趟探索之旅，我的感受是，将大型NLP模型推向移动端是一个充满挑战但非常有价值的方向。通过量化、剪枝和ONNX转换这套组合拳，我们确实能把模型的“体格”降下来，让它至少在技术上具备了在资源受限设备上运行的可能性。

这次尝试的 nlp_structbert_sentence-similarity_chinese-large 模型，即使经过压缩，对主流移动设备来说依然是个负担。它更像一个“技术可行性验证”，告诉我们极限在哪里。对于大多数真实的移动端应用，选择更小巧的模型架构，或者采用“云+端”协同的策略，可能是更务实的选择。

如果你正在为一个明确的、离线的、且对延迟要求不那么严苛的场景寻找解决方案，这套压缩和转换流程值得一试。建议先从量化开始，这是性价比最高的优化。如果效果和速度仍不满足，再考虑结合剪枝和微调，或者转向更小的模型。工具链（PyTorch, ONNX, ONNX Runtime）现在已经相当成熟，剩下的就是根据你的具体需求，在模型大小、推理速度和任务精度之间找到那个最佳的平衡点了。

---

获取更多AI镜像

想探索更多AI镜像和应用场景？访问 CSDN星图镜像广场，提供丰富的预置镜像，覆盖大模型推理、图像生成、视频生成、模型微调等多个领域，支持一键部署。