Chinese BERT 预训练模型下载与使用指南：从零开始避坑实践

1次阅读

没有评论

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

对于中文 NLP 开发者来说，下载和使用预训练模型是入门的第一步。然而，在实际操作中，往往会遇到以下几个典型问题：

网络延迟 ：直接从 Hugging Face 官网下载模型时，由于服务器位于海外，下载速度极慢，甚至经常中断。
库版本冲突 ：不同版本的 transformers 库对模型的加载方式可能存在差异，导致加载失败或运行时报错。
内存不足（OOM）：BERT 模型体积较大，尤其是在低配置机器上加载时，容易因内存不足而崩溃。
模型校验问题 ：下载过程中文件可能损坏，导致加载时出现校验错误。

这些问题往往会让新手开发者感到挫败，甚至放弃尝试。本文将针对这些问题，提供一站式的解决方案。

目前，下载 Chinese BERT 预训练模型主要有三种方式，各自优缺点如下：

使用 transformers 库的 from_pretrained 方法
优点：简单易用，自动处理模型缓存和版本兼容性。
缺点：依赖网络连接，下载速度慢，且无法断点续传。
使用镜像站下载
优点：国内镜像站（如清华源、阿里云）下载速度快，稳定性高。
缺点：需要手动配置镜像站地址，且部分镜像站可能更新不及时。
手动下载模型文件
优点：完全离线操作，适合网络环境极差的情况。
缺点：需要手动管理模型文件，容易出错。

综合来看，对于国内开发者，推荐优先使用镜像站下载，其次是 from_pretrained 方法，最后是手动下载。

1. 使用 `from_pretrained` 方法下载

以下是带异常处理和进度条的完整代码示例：

from transformers import BertModel, BertTokenizer
import os

# 配置模型缓存路径（避免默认路径权限问题）cache_dir = os.path.expanduser('~/.cache/huggingface')
os.makedirs(cache_dir, exist_ok=True)

# 下载模型
try:
    model = BertModel.from_pretrained('bert-base-chinese', cache_dir=cache_dir)
    tokenizer = BertTokenizer.from_pretrained('bert-base-chinese', cache_dir=cache_dir)
    print("模型下载成功！")
except Exception as e:
    print(f"模型下载失败: {e}")

可以通过设置环境变量 HF_ENDPOINT 来指定镜像站：

import os
from transformers import BertModel

# 使用清华镜像站
os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'

model = BertModel.from_pretrained('bert-base-chinese')

手动下载后，可以通过指定本地路径加载模型：

from transformers import BertModel

model = BertModel.from_pretrained('/path/to/bert-base-chinese')

BERT 模型占用内存较大，可以通过以下方式优化：

使用 FP16 精度加载模型 ：

model = BertModel.from_pretrained('bert-base-chinese', torch_dtype=torch.float16)

分批加载数据 ：避免一次性加载过多数据导致内存溢出。

确保 transformers 库和 torch 版本兼容：

import transformers
import torch

print(f"transformers 版本: {transformers.__version__}")
print(f"torch 版本: {torch.__version__}")

下载完成后，可以通过计算文件的哈希值来验证文件完整性：

import hashlib

def calculate_sha256(file_path):
    sha256_hash = hashlib.sha256()
    with open(file_path, "rb") as f:
        for byte_block in iter(lambda: f.read(4096), b""):
            sha256_hash.update(byte_block)
    return sha256_hash.hexdigest()

print(calculate_sha256("/path/to/model.bin"))

以下是不同下载方式的耗时对比（测试环境：国内网络，100Mbps 带宽）：