文本分类示例
step1 导入相关包
fromtransformersimportAutoTokenizer,AutoModelForSequenceClassification,Trainer,TrainingArgumentsfromdatasetsimportload_data_setstep2 加载数据集
dataset=load_dataset("csv",data_files="./ChnSentiCorp.csv",split="train")dataset=dataset.filter(lambdax:x["review"]isnotNone)step3 划分数据集
datasets=dataset.train_test_split(test_size=0.1)step4 数据集预处理
importtorch# 加载 rbt3 对应的分词器tokenizer=AutoTokenizer.from_pretrained("hfl/rbt3")# 定义预处理函数，batched=True时examples是一批样本defprocess_function(examples):# 对review字段（评论文本）分词编码tokenized_examples=tokenizer(examples["review"],max_length=128,truncation=True)# 把原始标签复制到labels字段（Trainer训练要求固定名字labels）tokenized_examples["labels"]=examples["label"]returntokenized_examples# 批量处理整个数据集tokenized_datasets=datasets.map(process_function,batched=True,remove_columns=datasets["train"].column_names)step5 创建模型
model=AutoModelForSequenceClassification.from_pretrained("hfl/chinese-macbert-large")step6 创建评估函数
importevaluate# 如果网络不太好，也可以使用本地加载的方式acc_metric=evaluate.load("./metric_accuracy.py")f1_metirc=evaluate.load("./metric_f1.py")defeval_metric(eval_predict):predictions,labels=eval_predict predictions=predictions.argmax(axis=-1)acc=acc_metric.compute(predictions=predictions,references=labels)f1=f1_metirc.compute(predictions=predictions,references=labels)acc.update(f1)returnaccStep7 创建TrainingArguments
train_args=TrainingArguments(output_dir="./checkpoints",# 输出文件夹per_device_train_batch_size=2,# 训练时的batch_sizegradient_accumulation_steps=32,# *** 梯度累加 ***gradient_checkpointing=True,# *** 梯度检查点 ***optim="adafactor",# *** adafactor优化器 ***per_device_eval_batch_size=4,# 验证时的batch_sizenum_train_epochs=1,# 训练轮数logging_steps=10,# log 打印的频率eval_strategy="epoch",# 评估策略save_strategy="epoch",# 保存策略save_total_limit=3,# 最大保存数learning_rate=2e-5,# 学习率weight_decay=0.001,# weight_decaymetric_for_best_model="f1",# 设定评估指标load_best_model_at_end=True)# 训练完成后加载最优模型train_argsStep8 创建Trainer
fromtransformersimportDataCollatorWithPadding# *** 参数冻结 ***forname,paraminmodel.bert.named_parameters():param.requires_grad=Falsetrainer=Trainer(model=model,args=train_args,tokenizer=tokenizer,train_dataset=tokenized_datasets["train"],eval_dataset=tokenized_datasets["test"],data_collator=DataCollatorWithPadding(tokenizer=tokenizer),compute_metrics=eval_metric)Step9 模型训练
trainer.train()trainer.evaluate(tokenized_datasets["test"])trainer.predict(tokenized_datasets["test"])Step10 模型预测
sen="我觉得这家酒店不错，饭很好吃！"id2_label={0:"差评！",1:"好评！"}model.eval()withtorch.inference_mode():inputs=tokenizer(sen,return_tensors="pt")inputs={k:v.cuda()fork,vininputs.items()}logits=model(**inputs).logits pred=torch.argmax(logits,dim=-1)print(f"输入：{sen}\n模型预测结果:{id2_label.get(pred.item())}")fromtransformersimportpipeline model.config.id2label=id2_label pipe=pipeline("text-classification",model=model,tokenizer=tokenizer,device=0)显存优化策略
模型权重 4B模型参数量
优化器状态 8B模型参数量，对于常用的AdamW优化器而言
梯度 4B*模型参数量
前向激活值 取决于序列长度、隐层维度、batch大小等多个因素
