在Mac M1/M2上使用Hugging Face Transformers进行中文文本分类(完整指南)
前言
随着Apple Silicon芯片(M1/M2)的普及,越来越多的开发者希望在Mac上运行深度学习任务。本文将详细介绍如何在Mac M1/M2设备上使用Hugging Face Transformers库进行中文文本分类任务,包括环境配置、数据处理、模型训练和性能优化等完整流程。
环境准备
1. 硬件和系统要求
• 设备:Apple M1/M2系列芯片的Mac
• 系统:macOS 12.3 (Monterey)或更高版本
• Python:3.8或更高版本
2. 安装必要的库
# 创建虚拟环境
python -m venv .venv
source .venv/bin/activate
# 安装支持MPS的PyTorch
pip install --pre torch torchvision --extra-index-url https://download.pytorch.org/whl/nightly/cpu
# 安装其他依赖
pip install transformers datasets evaluate
数据处理
1. 加载和预处理数据集
我们使用中文情感分析数据集ChnSentiCorp:
from datasets import load_from_disk
from transformers import AutoTokenizer
# 加载数据集
dataset = load_from_disk('./data/ChnSentiCorp')
# 缩小数据集规模
dataset['train'] = dataset['train'].shuffle().select(range(1500))
dataset['test'] = dataset['test'].shuffle().select(range(100))
# 初始化tokenizer
tokenizer = AutoTokenizer.from_pretrained('hfl/rbt3')
# 编码函数
def encode_data(data):
return tokenizer.batch_encode_plus(data['text'], truncation=True)
# 应用编码
dataset = dataset.map(encode_data, batched=True, batch_size=1000, num_proc=4, remove_columns=['text'])
# 过滤过长的句子
dataset = dataset.filter(lambda x: len(x['input_ids']) <= 512, batched=True)
2. 数据格式转换
from transformers import DataCollatorWithPadding
# 数据整理器
data_collator = DataCollatorWithPadding(tokenizer)
模型加载与配置
1. 加载预训练模型
from transformers import AutoModelForSequenceClassification
import torch
# 检查MPS是否可用
device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
print(f"Using device: {device}")
# 加载模型
model = AutoModelForSequenceClassification.from_pretrained('./model/rbt3', num_labels=2)
model = model.to(device)
2. 自定义训练器(适配MPS)
from transformers import Trainer
class MPSReadyTrainer(Trainer):
def prediction_step(self, model, inputs, prediction_loss_only=False, ignore_keys=None):
inputs = {k: v.to('mps') for k, v in inputs.items()}
return super().prediction_step(model, inputs, prediction_loss_only, ignore_keys)
训练配置
1. 设置训练参数
from transformers import TrainingArguments
training_args = TrainingArguments(
output_dir='./out/rbt3',
evaluation_strategy="steps",
eval_steps=30,
save_strategy='steps',
save_steps=30,
learning_rate=5e-5,
per_device_train_batch_size=16, # M1/M2建议较小batch size
per_device_eval_batch_size=16,
num_train_epochs=10,
logging_dir='./log/rbt3',
no_cuda=True,
use_mps_device=True,
fp16=False # MPS暂不支持混合精度
)
2. 定义评估指标
from evaluate import load
def compute_metrics(eval_pred):
metric = load('accuracy')
logits, labels = eval_pred
if isinstance(logits, torch.Tensor):
predictions = logits.argmax(dim=-1)
else:
predictions = torch.from_numpy(logits).argmax(dim=-1)
return metric.compute(predictions=predictions, references=labels)
训练与评估
1. 初始化训练器
trainer = MPSReadyTrainer(
model=model,
args=training_args,
train_dataset=dataset['train'],
eval_dataset=dataset['test'],
compute_metrics=compute_metrics,
data_collator=data_collator
)
2. 开始训练
print("========== 开始训练 ==========")
trainer.train()
print("========== 最终评估 ==========")
trainer.evaluate()
性能优化技巧
- 调整batch size:M1/M2芯片建议使用8-32的batch size
- 禁用pin_memory:在TrainingArguments中设置
dataloader_pin_memory=False
- 减少数据加载线程:设置
num_proc=2
或更低 - 简化模型:使用更小的预训练模型如’rbt3’而非’bert-base’
常见问题解决
MPS不可用错误
• 确保安装了正确版本的PyTorch
• 检查macOS版本≥12.3
• 运行python -c "import torch; print(torch.backends.mps.is_available())"
确认内存不足错误
• 减小batch size
• 缩短序列长度(max_length=256)数据类型不匹配
• 确保所有张量都通过.to(device)
转移到MPS
结语
本文详细介绍了在Apple Silicon Mac上使用Hugging Face Transformers进行中文文本分类的完整流程。通过合理配置和优化,可以在Mac设备上高效地进行NLP模型训练。希望这篇指南能帮助开发者充分利用M1/M2芯片的性能优势。
完整代码已上传GitHub:项目链接
问题讨论欢迎在评论区留言