Android端侧可部署的小语言模型实战：从模型选型到性能优化

1次阅读

没有评论

共计 2832 个字符，预计需要花费 8 分钟才能阅读完成。

移动设备上运行 NLP 任务时，开发者常遇到三个致命问题：

内存瓶颈：BERT-base 模型动辄 400MB+ 内存占用，低端设备直接 OOM
延迟敏感：云端 API 的网络往返时间（通常 200-500ms）难以满足实时交互需求
隐私顾虑：用户对话记录、输入内容上传云端存在合规风险

DistilBERT：BERT 的蒸馏版，参数量减少 40% 但保留 97% 性能
TinyLLaMA：专为移动端优化的 1.1B 参数模型，支持中英双语
MobileBERT：谷歌针对移动设备优化的版本，内置分组注意力机制

TensorFlow Lite：支持量化、GPU 代理、XNNPACK 加速，灵活度高
ML Kit：谷歌全家桶方案，但自定义模型能力有限
ONNX Runtime：跨平台优势明显，适合多端统一部署

我们最终选择 DistilBERT + TFLite 组合，平衡了性能与灵活性。

# 使用 transformers 和 onnxruntime 工具链
from transformers import DistilBertTokenizer, DistilBertModel
import torch
import onnx
from onnxruntime.quantization import quantize_dynamic, QuantType

# 原始模型导出 ONNX
model = DistilBertModel.from_pretrained('distilbert-base-uncased')
dummy_input = torch.zeros(1, 128, dtype=torch.long)
torch.onnx.export(model, dummy_input, "distilbert.onnx")

# 动态量化（8-bit 整数）quantize_dynamic(
    "distilbert.onnx",
    "distilbert_int8.onnx",
    weight_type=QuantType.QInt8
)

量化后模型大小从 255MB 降至 68MB，内存占用减少约 65%。

build.gradle 关键配置：

android {
    defaultConfig {
        ndk {abiFilters 'armeabi-v7a', 'arm64-v8a'}
    }
    externalNativeBuild {
        cmake {
            arguments "-DANDROID_TOOLCHAIN=clang"
            cppFlags "-march=armv8.2-a+dotprod"  # 启用 ARM DSP 指令集
        }
    }
}

JNI 接口示例：

#include <jni.h>
#include "tensorflow/lite/interpreter.h"

extern "C" JNIEXPORT jfloatArray JNICALL
Java_com_example_nlp_NativeLib_runInference(
    JNIEnv* env,
    jobject /* this */,
    jintArray input_ids) {

    // 1. 获取输入缓冲区
    TfLiteTensor* input = interpreter->input_tensor(0);
    int* in_data = input->data.i32;

    // 2. 执行推理
    interpreter->Invoke();

    // 3. 返回输出
    TfLiteTensor* output = interpreter->output_tensor(0);
    jfloatArray result = env->NewFloatArray(output->bytes);
    env->SetFloatArrayRegion(result, 0, output->bytes, output->data.f);
    return result;
}

模型分片加载：将模型拆分为多个.tflite 文件，按需加载
内存映射 ：使用MappedByteBuffer 避免完整模型载入内存
动态卸载 ：在 onTrimMemory() 回调中释放非活跃模型

class ModelLoader(context: Context) {private val modelMap = mutableMapOf<String, Interpreter>()

    fun loadModel(name: String): Interpreter {return modelMap.getOrPut(name) {val assetFileDescriptor = context.assets.openFd("$name.tflite")
            val inputStream = FileInputStream(assetFileDescriptor.fileDescriptor)
            val modelBuffer = inputStream.channel.map(
                FileChannel.MapMode.READ_ONLY,
                assetFileDescriptor.startOffset,
                assetFileDescriptor.declaredLength
            )
            Interpreter(modelBuffer)
        }
    }

    fun trimMemory(level: Int) {if (level >= ComponentCallbacks2.TRIM_MEMORY_RUNNING_CRITICAL) {modelMap.clear()
        }
    }
}

测试设备：Redmi Note 10 Pro（骁龙 732G）