前置知識(shí)
- Colab 的使用
- Huggingface 官網(wǎng)和一些基礎(chǔ)API
上代碼
- 首先,建議保存代碼到
VSCode
,這樣雙擊關(guān)鍵類(lèi),F12
可以進(jìn)入查看具體接口參數(shù)及其含義。
然后,建議代碼在Colab
運(yùn)行,第一個(gè)是有默認(rèn)GPU
資源,第二個(gè)是否則會(huì)產(chǎn)生各種ConnectionError, OSError
等錯(cuò)誤… - 重點(diǎn)可以看注釋。自行摸索了一些額外的參數(shù),大多數(shù)人都沒(méi)有講訓(xùn)練中/后需要保存模型參數(shù)…
"""
首先運(yùn)行如下代碼安裝庫(kù)
然后直接運(yùn)行改代碼即可
!pip install datasets transformers
!pip install accelerate -U
"""
from datasets import load_dataset
from transformers import (
AutoTokenizer,
DataCollatorWithPadding,
TrainingArguments,
AutoModelForSequenceClassification,
Trainer,
)
# 加載數(shù)據(jù)集,并加載對(duì)應(yīng)模型的分詞器
raw_datasets = load_dataset("glue", "mrpc")
checkpoint = "bert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(checkpoint)
def tokenize_function(example):
return tokenizer(example["sentence1"], example["sentence2"], truncation=True)
# 數(shù)據(jù)集分詞并打包,傳給data_collator
tokenized_datasets = raw_datasets.map(tokenize_function, batched=True)
data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
# 設(shè)置訓(xùn)練參數(shù),這里我選擇訓(xùn)練1poch,每處理20%steps就保存,注意最后100%時(shí)不保存。
training_args = TrainingArguments(
"test-trainer",
num_train_epochs=1,
save_strategy="steps",
save_steps=0.2,
)
# 設(shè)置模型
model = AutoModelForSequenceClassification.from_pretrained(checkpoint, num_labels=2)
# 設(shè)置訓(xùn)練器,提供各種必要參數(shù)。
trainer = Trainer(
model,
training_args,
train_dataset=tokenized_datasets["train"],
eval_dataset=tokenized_datasets["validation"],
data_collator=data_collator,
tokenizer=tokenizer,
)
# 訓(xùn)練,結(jié)束后保存模型
trainer.train()
model.save_pretrained("./output_model")
- 最后文件夾如下,test-trainer 保存訓(xùn)練斷點(diǎn),output_model保存訓(xùn)練后參數(shù)模型。
文章來(lái)源地址http://www.zghlxwxcb.cn/news/detail-544405.html
文章來(lái)源:http://www.zghlxwxcb.cn/news/detail-544405.html
到了這里,關(guān)于【NLP,Huggingface,Colab】使用 Trainer 訓(xùn)練模型,并保存模型參數(shù)的文章就介紹完了。如果您還想了解更多內(nèi)容,請(qǐng)?jiān)谟疑辖撬阉鱐OY模板網(wǎng)以前的文章或繼續(xù)瀏覽下面的相關(guān)文章,希望大家以后多多支持TOY模板網(wǎng)!