中文BERT预训练模型下载与部署实战指南：从Hugging Face到生产环境

1次阅读

没有评论

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

最近在做一个中文文本分类项目时，我发现部署 BERT 模型这件事，远比想象中复杂。主要遇到三个问题：

下载速度慢：从 Hugging Face 官方源下载 bert-base-chinese 模型，速度经常只有几十 KB/s，一个 600MB 的模型要下几个小时
版本混乱：光是中文 BERT 就有 bert-base-chinese、chinese-bert-wwm、chinese-roberta-wwm-ext 等多个变体，不知道选哪个好
部署困难：直接加载原始模型显存占用太高，在消费级 GPU 上根本跑不起来

经过对比测试，国内常用的中文 BERT 变体主要有：

bert-base-chinese：最基础的官方中文版，适合大多数任务
chinese-bert-wwm：使用全词掩码 (Whole Word Masking) 技术，对中文更友好
chinese-roberta-wwm-ext：哈工大发布的 RoBERTa 版，在多个中文任务上表现更好

我最终选择了chinese-roberta-wwm-ext，因为它在 CLUE 榜单上的表现最好。

直接从 Hugging Face 下载太慢，我找到了两个替代方案：

清华大学镜像源：

from transformers import BertTokenizer, BertModel

tokenizer = BertTokenizer.from_pretrained('https://mirrors.tuna.tsinghua.edu.cn/hugging-face-models/bert-base-chinese')
model = BertModel.from_pretrained('https://mirrors.tuna.tsinghua.edu.cn/hugging-face-models/bert-base-chinese')

阿里云 OSS 备份：

wget https://bert-models.oss-cn-beijing.aliyuncs.com/chinese-roberta-wwm-ext/pytorch_model.bin

实测下载速度从 50KB/ s 提升到 8MB/s，600MB 的模型 1 分钟就能下完。

原始 BERT 模型太大，我用 ONNX 量化将模型压缩了 4 倍：

from transformers import BertModel
import torch

# 加载原始模型
model = BertModel.from_pretrained('chinese-roberta-wwm-ext')

# 转换为 ONNX 格式
dummy_input = torch.randn(1, 128, dtype=torch.long)
torch.onnx.export(model, dummy_input, "model.onnx")

# 量化模型
import onnxruntime as ort
from onnxruntime.quantization import quantize_dynamic

quantize_dynamic("model.onnx", "model_quant.onnx")

量化前后对比：

指标	原始模型	量化模型
大小	438MB	112MB
显存占用	1.2GB	320MB
推理速度(ms)	45	22

遇到过最头疼的问题就是 CUDA 版本不匹配。我的经验是：

先用 nvidia-smi 查看驱动支持的最高 CUDA 版本
根据 PyTorch 官网的兼容性表格选择对应版本
用 conda 安装指定版本的 CUDA 工具包：
```
conda install cudatoolkit=11.1
```

当显存不够时，可以使用分块加载：

from transformers import BertConfig, BertModel

# 先加载配置
config = BertConfig.from_pretrained('chinese-roberta-wwm-ext')

# 再分块加载模型
model = BertModel.from_pretrained('chinese-roberta-wwm-ext', 
                                config=config,
                                device_map='auto',
                                low_cpu_mem_usage=True)

为了避免环境问题，我推荐使用 Docker 部署：

FROM pytorch/pytorch:1.11.0-cuda11.3-cudnn8-runtime

RUN pip install transformers==4.21.0 onnxruntime-gpu

COPY . /app
WORKDIR /app

CMD ["python", "app.py"]

在多线程环境下，建议这样加载模型：

import threading
from transformers import BertModel

model_lock = threading.Lock()
model = None

def get_model():
    global model
    if model is None:
        with model_lock:
            if model is None:
                model = BertModel.from_pretrained('chinese-roberta-wwm-ext')
    return model

我现在的做法是：