huggingface NLP工具包教程3：微調預訓練模型

引言

在上一章我們已經介紹了如何使用 tokenizer 以及如何使用預訓練的模型來進行預測。本章將介紹如何在自己的數據集上微調一個預訓練的模型。在本章，你將學到：

• 如何從 Hub 準備大型數據集
• 如何使用高層 Trainer API 微調模型
• 如何使用自定義訓練循環
• 如何利用 Accelerate 庫，進行分布式訓練

如果想要將將經過訓練的權重上傳到 Hugging Face Hub，需要注冊一個 huggingface.co 賬號：創建賬號。

處理數據

以下是如何訓練一個序列分類器（以一個 batch 為例）：

import torch
from transformers import AdamW, AutoTokenizer, AutoModelForSequenceClassification# 與之前章節一致
checkpoint = "bert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(checkpoint)
model = AutoModelForSequenceClassification.from_pretrained(checkpoint)
sequences = ["I've been waiting for a HuggingFace course my whole life.","This course is amazing!",
]
batch = tokenizer(sequences, padding=True, truncation=True, return_tensors="pt")# 以下是訓練部分
batch["labels"] = torch.tensor([1, 1])optimizer = AdamW(model.parameters())
loss = model(**batch).loss
loss.backward()
optimizer.step()

當然，向上面這樣僅使用兩個句子進行訓練肯定無法得到滿意的結果。我們需要一個更大的數據集。

本節我們將以 MRPC (Microsoft Research Paraphrase Corpus) 數據集為例，該數據集由 5801 對句子組成，并帶有一個標簽，表明它們是否是轉述（即，如果兩個句子的意思相同）。本章選擇該數據集，因為它是一個小數據集，所以很容易進行訓練。

從Hub中加載數據

Hub 不僅包括模型，它同樣包含不同語言的多個數據集。可以在這里查看數據集，推薦讀者在過完本節之后自己試著加載并處理一個新數據集，文檔參考這里。不過現在，讓我們先來看 MRPC 數據集，他是組成 GLUE benchmark 的十個數據集之一。GLUE benchmark 是一個在 10 個不同的文本分類任務上評估機器學習模型的學術基準。

Datasets 庫提供了一些非常簡單的命令來下載并緩存 Hub 中的數據集。例如下載 MRPC 數據集：

from datasets import load_datasetraw_datasets = load_dataset("glue", "mrpc")
print(raw_datasets)# 輸出：
DatasetDict({train: Dataset({features: ['sentence1', 'sentence2', 'label', 'idx'],num_rows: 3668})validation: Dataset({features: ['sentence1', 'sentence2', 'label', 'idx'],num_rows: 408})test: Dataset({features: ['sentence1', 'sentence2', 'label', 'idx'],num_rows: 1725})
})

可以看到，我們得到了一個 DatasetDict 對象，其中包含了訓練集、驗證集和測試集。其中每個又包括了幾個列（sentence1,sentence2,labe,idx），和一個行數值，表示每個集合中的樣本個數。

上述命令會下載并緩存指定數據集，默認緩存目錄是 ~/.cache/huggingface/datasets。同樣可以通過環境變量 HF_HOME 來修改。

我們可以通過索引來訪問 raw_datasets 中的每對句子：

raw_train_dataset = raw_datasets["train"]
print(raw_train_dataset[0])# 輸出：
{'idx': 0,'label': 1,'sentence1': 'Amrozi accused his brother , whom he called " the witness " , of deliberately distorting his evidence .','sentence2': 'Referring to him as only " the witness " , Amrozi accused his brother of deliberately distorting his evidence .'}

可以看到標簽已經是整型了，所以這里不再需要進一步處理。如果想要知道哪個整型值對應哪個標簽，可以查看 raw_train_dataset 的 features 屬性。它會返回每一列的類型：

print(raw_train_dataset.features)# 輸出：
{'sentence1': Value(dtype='string', id=None),'sentence2': Value(dtype='string', id=None),'label': ClassLabel(num_classes=2, names=['not_equivalent', 'equivalent'], names_file=None, id=None),'idx': Value(dtype='int32', id=None)}

可以看到，標簽的類型為 ClassLabel，整數到標簽名稱的映射存儲在 names folder 中。0 對應于 not_equivalent，1 對應于 equivalent。

數據集預處理

我們需要將原文本數據轉換為模型可以接收的數值型，之前的章節已經介紹過，這由 tokenizer 完成。tokenizer 既可以接收一個句子，也可以接收一組句子。因此，我們可以直接將數據集中每個句子對中的 ”第一個句子“ 和 ”第二個句子” 直接送入給 tokenizer：

from transformers import AutoTokenizercheckpoint = "bert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(checkpoint)
tokenized_sentences_1 = tokenizer(raw_datasets["train"]["sentence1"])
tokenized_sentences_2 = tokenizer(raw_datasets["train"]["sentence2"])

然而，我們不能僅僅是直接將兩個序列傳給模型，然后讓模型預測兩個序列是否是釋義關系。我們需要將兩個序列處理成一個序列對，并進行適當的預處理。幸運的是，強大的 tokenizer 也可以接收一對序列并將它們處理成 BERT 模型需要的形式：

inputs = tokenizer("This is the first sentence.", "This is the second one.")
print(inputs)# 輸出：
{ 'input_ids': [101, 2023, 2003, 1996, 2034, 6251, 1012, 102, 2023, 2003, 1996, 2117, 2028, 1012, 102],'token_type_ids': [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1],'attention_mask': [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]
}

輸出中的 inpud_ids 和 attention_mask 我們在前一章已經介紹過，但是 token_type_ids 當時沒有提，在這里，它負責告訴模型哪部分輸入是第一個句子，哪部分是第二個句子。

我們可以對 input_ids 進行解碼會單詞：

tokenizer.convert_ids_to_tokens(inputs["input_ids"])# 輸出：
['[CLS]', 'this', 'is', 'the', 'first', 'sentence', '.', '[SEP]', 'this', 'is', 'the', 'second', 'one', '.', '[SEP]']

可以看到，模型期望的輸入是將兩個句子合并成 [CLS] sentence1 [SEP] sentence2 [SEP] 的形式，這與我們給出的 token_type_ids 是對齊的：

['[CLS]', 'this', 'is', 'the', 'first', 'sentence', '.', '[SEP]', 'this', 'is', 'the', 'second', 'one', '.', '[SEP]']
[      0,      0,    0,     0,       0,          0,   0,       0,      1,    1,     1,        1,     1,   1,       1]

第一個句子 [CLS] sentence1 [SEP] 對應的 token_type_ids 都為 0，而第二個句子 sentence2 [SEP] 對應的 token_type_ids 為 1。

注意 token_id_types 并不是所有模型都必須的，只有模型預訓練任務中需要這種輸入時才需要。也就是說如果我們用的是其他預訓練模型（比如 DistilBERT），可能就不需要 token_id_types，這時 tokenizer 也不會返回該鍵。

這里的 BERT 模型預訓練時是需要 token_type_ids 的，BERT 模型在預訓練時，除了第一章中介紹過的 masked language modeling 任務之外，還有另一個預訓練任務：next sentence prediction ，該任務的目標是建模一對句子之間的關系。具體來說，該任務需要預測一對句子（當然也包括一些 mask 掉的 token）之中，第二句在語義上是否是第一句的后一句。為了使得該任務有一定難度，一半訓練數據中兩個句子來自同一個文本（是或不是連續兩句），另一半訓練數據則來自不同文本（肯定不是連續兩句）。

總之，只要保證加載的模型和加載的 tokenizer 是來自同一個預訓練權重，我們就不太需要擔心輸入數據中是否需要包含 token_type_ids ，因為 tokenizer 會幫我們把一切都準備好。

我們已經了解如何使用 tokenizer 來處理一對句子，現在用它來處理整個數據集：直接將一組句子對送入到 tokenizer 中，并且指定填充和截斷：

tokenized_dataset = tokenizer(raw_datasets["train"]["sentence1"],raw_datasets["train"]["sentence2"],padding=True,truncation=True,
)

這樣可以正常工作，這樣的缺點是它會直接返回一個字典（包括 input_ids, attention_mask, and token_type_ids），這需要我們的機器有足夠大的內存來存儲整個數據集。

我們會使用 Dataset.map() 方法來封裝預處理過程，該方法比較靈活，方便除了 tokenization 之外，在預處理階段添加更多操作。map() 方法的工作方式是將一個函數內定義的操作施加到 Dataset 中的每個元素上。這里我們先定義一個處理函數：

def tokenize_function(example):return tokenizer(example["sentence1"], example["sentence2"], truncation=True)

該函數接收一個字典（比如數據集中的一個樣本）作為輸入，并返回一個新的字典，包含 input_ids, attention_mask, and token_type_ids 這幾個鍵。如果 example 字典包含多個樣本（每個鍵都是一個句子列表），也是可行的。因為前面介紹過，tokenizer 可以處理成對的句子列表。我們在調用 map() 的時傳入參數 batched=True，這將大大加快 tokenization 過程。tokenizer 來自由 Rust 編寫的 Tokenizer 庫。

注意我們并沒有在 tokenize_function 中設置 padding 參數，因為直接按照全數據集的最大長度進行填充是很低效的，最好是在構建 batch 時進行填充，這樣就只需要按照 batch 內的最大長度進行填充即可，而不需要按照整個數據集的最大長度進行填充。

我們通過設置 batched=True 來同時處理多個元素，而非一個一個處理，這會使得整個預處理過程更快：

tokenized_datasets = raw_datasets.map(tokenize_function, batched=True)
print(tokenized_datasets)# 輸出：
DatasetDict({train: Dataset({features: ['attention_mask', 'idx', 'input_ids', 'label', 'sentence1', 'sentence2', 'token_type_ids'],num_rows: 3668})validation: Dataset({features: ['attention_mask', 'idx', 'input_ids', 'label', 'sentence1', 'sentence2', 'token_type_ids'],num_rows: 408})test: Dataset({features: ['attention_mask', 'idx', 'input_ids', 'label', 'sentence1', 'sentence2', 'token_type_ids'],num_rows: 1725})
})

可以看到，map 會通過在字典中添加處理后得到的新鍵來完成預處理過程。如果某些處理結果會得到已有鍵的新值，map 也會用新值覆蓋掉舊值。

我們還可以在調用 map 方法時，通過傳入 num_proc 參數來進行多進程的數據預處理。上例中沒有這么做是因為 tokenizer 庫已經使用了多線程來進行加速，但如果在實際中沒有使用該庫的 tokenizer 的話，可以在這里使用多進程進行加速。

最后一件要介紹的事情是動態填充（dynamic padding），它負責將一個 batch 內的所有序列填充到該 batch 內原始序列的最長長度。

動態填充

collate function 負責將一個 batch 內所有的樣本放到一起。在 Pytorch 中，它是我們構建 DataLoader 時一個可選的參數，默認的 collate function 會簡單地將所有的樣本數據轉換為張量并拼接在一起。這肯定是不能直接用的，因為目前我們還沒有進行填充，batch 內的每個樣本不是等長的。之前我們故意先沒有填充，因為我們想在構建每個 batch 時進行填充，避免一個 batch 內無意義的填充。這將大大加快訓練速度，但請注意，如果時在TPU上訓練，可能會導致問題，因為 TPU 更偏好固定的形狀，即使這需要額外的填充。

實際中，我們會定義一個 collate function，來對每個 batch 內的樣本進行填充。Transformers 庫通過 DataCollatorWithPadding 提供這樣的功能。在對其進行實例化時，傳入我們的 tokenizer，因為它需要知道 padding token 是什么，并且要知道該模型需要再輸入序列的左側填充還是右側填充，之后，它會為我們做好一切。

from transformers import DataCollatorWithPaddingdata_collator = DataCollatorWithPadding(tokenizer=tokenizer)

我們從我們的訓練集中抓取一些樣本來將它們打包成 batch。這里，我們刪除列idx、sentence1 和 sentence2，因為不需要它們而且它們包含字符串（我們不能用字符串創建張量），并查看批次中每個元素的長度：

samples = tokenized_datasets["train"][:8]
samples = {k: v for k, v in samples.items() if k not in ["idx", "sentence1", "sentence2"]}
[len(x) for x in samples["input_ids"]]# 輸出：
[50, 59, 47, 67, 59, 50, 62, 32]

毫無疑問，一個批次中的原始序列長度是不一致的，從 32 到 67。動態填充意味著該批次內的序列都要填充到序列中的最長長度為 67。如果不使用動態填充的話，這些樣本的序列長度都要被填充到整個數據集的最大長度，或者模型能夠接收的最大長度。下面我們再檢查一下 data_collator 是否正確進行了動態填充：

batch = data_collator(samples)
{k: v.shape for k, v in batch.items()}# 輸出：
{'attention_mask': torch.Size([8, 67]),'input_ids': torch.Size([8, 67]),'token_type_ids': torch.Size([8, 67]),'labels': torch.Size([8])}

結果一切正常。至此，我們就完成了對原始文本數據的預處理，準備正式開始進行微調。

使用Trainer API對模型進行微調

Transformers 庫提供了一個 Trainer 類來幫助用戶在自己的數據集上對預訓練模型進行微調。在完成上一節的數據預處理準備之后，馬上就可以開始微調訓練了。

下面的代碼是上一節數據預處理的匯總：

from datasets import load_dataset
from transformers import AutoTokenizer, DataCollatorWithPaddingraw_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)tokenized_datasets = raw_datasets.map(tokenize_function, batched=True)
data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

訓練

在我們定義一個 Trainer 類之前，第一步要做的是定義一個 TrainingArguments 類，其中包括了 Trainer 訓練和驗證時所需的所有超參數。我們唯一必須要提供的參數時模型和權重參數的存放目錄，其他的參數均默認，對于一個基礎的微調訓練，這樣就可以工作。

from transformers import TrainingArgumentstraining_args = TrainingArguments("test-trainer")

第二步就是要定義模型。這里我們使用 AutoModelForSequenceClassification，類別數為 2：

from transformers import AutoModelForSequenceClassificationmodel = AutoModelForSequenceClassification.from_pretrained(checkpoint, num_labels=2)

這一步實例化一個模型之后，會得到一個警告。這是因為 BERT 模型預訓練任務不是對句子對進行分類，所以預訓練時的模型頭部被直接丟掉，然后換用一個新的適合指定任務的頭部。該警告表明預訓練模型中有一部分權重（此處就是模型頭部部分）沒有用到，并且另外有一些權重（此處即新的頭部）是隨機初始化的。

當我們定義好模型之后，就可以定義 Trainer 了，將我們目前得到的對象都丟進去：

from transformers import Trainertrainer = Trainer(model,training_args,train_dataset=tokenized_datasets["train"],eval_dataset=tokenized_datasets["validation"],data_collator=data_collator,tokenizer=tokenizer,
)

按照上述方式傳入 tokenizer 之后，trainer 使用的 data_collator 將會是我們之前定義的 DataCollatorWithPadding ，所以實際上 data_collator=data_collator 這一行是可以跳過的。

接下來，直接調用 trainer.train() 方法就可以開始微調模型：

trainer.train()

這就會開始微調，并每過 500 個 steps 就報告一次損失。但是這并不能告訴我們模型實際性能如何，因為：

1. 我們沒有通過設置 evaluation_strategy 來告訴模型在每個 step 或每個 epoch 之后對模型進行評估
2. 我們沒有提供 compute_metrics() 函數給模型，來告訴他如何計算指標

evaluation

接下來介紹如何構建一個有用的 compute_metrics() 函數，并在訓練時使用它。該函數必須接受一個 EvalPrediction 對象（它是一個帶 predictions 字段和 label_ids 字段的命名元組），并將返回一個字典，將字符串映射為浮點值（字符串是返回的指標的名字和浮點值結果）。為了獲得模型的預測結果，我們可以使用 Trainer.predict() 方法：

predictions = trainer.predict(tokenized_datasets["validation"])
print(predictions.predictions.shape, predictions.label_ids.shape)# 輸出：
(408, 2) (408,)

predict() 方法的返回值是一個命名元組，共有三個字段：predictions, label_ids, metrics 。metrics 字段僅包含損失值和一些時間指標（預測的總時長和平均時長）。當我們寫好 compute_metrics() 函數并將他傳給 trainer 之后，該字段的返回值就會包括 compute_metrics() 返回的指標。

可以看到 predictions 是一個二維數組，形狀為 $408\times 2$ （ 408 是數據集中的樣本個數）。這是數據集中每個樣本預測結果 logits，我們需要取它們最大值的索引，來得到模型最終預測的類別：

import numpy as nppreds = np.argmax(predictions.predictions, axis=-1)

得到最終的預測類別之后，就可以與標簽進行對比，計算指標。我們基于 Evaluate 庫來構建 compute_metric() 函數。加載 MRPC 數據集的相關指標，與加載數據集一樣簡單，這次我們使用 evaluate.load() 函數，它會返回一個對象，該對象有 compute() 方法，我們可以直接用來計算指標：

import evaluatemetric = evaluate.load("glue", "mrpc")
metric.compute(predictions=preds, references=predictions.label_ids)# 輸出：
{'accuracy': 0.8578431372549019, 'f1': 0.8996539792387542}

由于模型頭部是隨機初始化的，因此最終的測試結果數值可能會稍有不同。這里我們看到模型在驗證集上的準確率為 85.78%，F1 分數為 89.97。這是 GLUE Benchmark 上評測 MRPC 數據集所用的指標。在 BERT 原論文中報告的結果中，base 模型的 F1 分數為 88.9。論文中使用的是 uncased 模型，這里我們用的是 cased 模型，因此結果稍好。

將所有東西封裝在一起，就得到了我們的 compute_metrics() 函數：

def compute_metrics(eval_preds):metric = evaluate.load("glue", "mrpc")logits, labels = eval_predspredictions = np.argmax(logits, axis=-1)return metric.compute(predictions=predictions, references=labels)

為了在每一個 epoch 結束時查看這些指標，我們重新定義一個 Trainer，將 compute_metrics 函數加進來：

training_args = TrainingArguments("test-trainer", evaluation_strategy="epoch")
model = AutoModelForSequenceClassification.from_pretrained(checkpoint, num_labels=2)trainer = Trainer(model,training_args,train_dataset=tokenized_datasets["train"],eval_dataset=tokenized_datasets["validation"],data_collator=data_collator,tokenizer=tokenizer,compute_metrics=compute_metrics,
)

我們在新的 TrainingArguments 中加入了 evaluation='epoch'，并創建了新的模型。啟動訓練：

trainer.train()

這一次，它將在每個 epoch 結束時，除了會報告訓練損失，還會報告驗證損失和和指標。同樣，由于模型的隨機頭部初始化，這次達到的準確率和F1分數可能與之前有所不同，但差別不會太大。

Trainer 可在多個 GPU 或 TPU 上開箱即用，并提供許多選項，如混合精度訓練（在訓練參數中使用 fp16=True）。我們將在第10章介紹更多。

使用Trainer API進行微調的介紹到此結束。第7章將給出為最常見的NLP任務執行此操作的示例，現在讓我們先看看如何在純 PyTorch 中執行相同的操作。

完整訓練

現在我們將看到如何在不使用 Trainer 類，而使用純 Pytorch，做到與上一節相同的事情。再次回顧第二節數據處理如下：

from datasets import load_dataset
from transformers import AutoTokenizer, DataCollatorWithPaddingraw_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)tokenized_datasets = raw_datasets.map(tokenize_function, batched=True)
data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

準備訓練

在實際編寫訓練腳本之前，我們需要先定義幾個對象。第一個是我們將用于迭代batch 數據加載程序。但在定義這些 DataLoader 之前，我們需要對經過 tokenize 的數據集進行一些后處理，在上一節中 Trainer 自動為我們做了這些事情。具體來說，我們需要：

• 刪除與模型不需要的值相對應的列（如句子1和句子2列）。
• 將列 label 重命名為 labels（因為模型需要的對應參數名為 labels）。
• 設置數據集的格式，使得它們返回 PyTorch 張量而不是列表。

tokenized_dataset 為每個步驟提供了對應方法：

tokenized_datasets = tokenized_datasets.remove_columns(["sentence1", "sentence2", "idx"])
tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
tokenized_datasets.set_format("torch")
tokenized_datasets["train"].column_names

然后檢查一下結果是否對應我們模型需要的鍵：

["attention_mask", "input_ids", "labels", "token_type_ids"]

現在準備工作做好了，開始定義 dataloader：

from torch.utils.data import DataLoadertrain_dataloader = DataLoader(tokenized_datasets["train"], shuffle=True, batch_size=8, collate_fn=data_collator
)
eval_dataloader = DataLoader(tokenized_datasets["validation"], batch_size=8, collate_fn=data_collator
)

可以用過以下方法快速檢查數據加載過程沒有錯誤：

for batch in train_dataloader:break
print({k: v.shape for k, v in batch.items()})# 輸出：
{'attention_mask': torch.Size([8, 65]),'input_ids': torch.Size([8, 65]),'labels': torch.Size([8]),'token_type_ids': torch.Size([8, 65])}

訓練數據的 Dataloader 設置了 shuffle=True，并且在 batch 中填充了最大長度，因此每個人實際查看的形狀可能會略有不同。

現在已經完全完成了數據預處理，開始準備模型。我們與上一節中所做的完全一樣進行實例化：

from transformers import AutoModelForSequenceClassificationmodel = AutoModelForSequenceClassification.from_pretrained(checkpoint, num_labels=2)

傳入一個 batch 測試有無問題：

outputs = model(**batch)
print(outputs.loss, outputs.logits.shape)# 輸出：
tensor(0.5441, grad_fn=<NllLossBackward>) torch.Size([8, 2])

當提供標簽時，所有 Transformer 模型都會返回損失，也會得到 logits（在我們的 batch 中，每次輸入兩個，所以張量大小為 $8x2$）。

還差兩個東西：優化器（optimizer）和學習率調度器（learning rate scheduler）。由于我們試圖以手動操作復現 trainer 的結果，因此這里使用相同的默認值。Trainer 使用的優化器是AdamW，它與Adam 相同，但對權重衰減正則項進行了扭曲（參見“Decoupled Weight Decay Regularization” ）：

from transformers import AdamWoptimizer = AdamW(model.parameters(), lr=5e-5)

最后，默認情況下使用的學習速率調度器是從最大值（5e-5）到 0 的線性衰減。我們需要知道總訓練步數，即我們要運行的 epoch 數乘以訓練 batch 數（即 DataLoader 的長度）。trainer 默認使用三個 epoch，因此我們按照如下定義：

from transformers import get_schedulernum_epochs = 3
num_training_steps = num_epochs * len(train_dataloader)
lr_scheduler = get_scheduler("linear",optimizer=optimizer,num_warmup_steps=0,num_training_steps=num_training_steps,
)
print(num_training_steps)# 輸出：
1377

訓練循環

然后我們再定義一下訓練所用的設備：

import torchdevice = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
model.to(device)
print(device)# 輸出：
device(type='cuda')

現在可以編寫訓練腳本并進行訓練了：

from tqdm.auto import tqdmprogress_bar = tqdm(range(num_training_steps))model.train()
for epoch in range(num_epochs):for batch in train_dataloader:batch = {k: v.to(device) for k, v in batch.items()}outputs = model(**batch)loss = outputs.lossloss.backward()optimizer.step()lr_scheduler.step()optimizer.zero_grad()progress_bar.update(1)

可以看到，訓練循環的核心步驟看起來很像引言中。我們沒有打印任何指標，所以這個訓練循環不會告訴我們關于模型如何運行的任何信息。我們需要為此添加一個評估循環。

評估循環

與前面一樣，我們將使用 Evaluate 庫。我們介紹過了 metric.compute() 方法：

import evaluatemetric = evaluate.load("glue", "mrpc")
model.eval()
for batch in eval_dataloader:batch = {k: v.to(device) for k, v in batch.items()}with torch.no_grad():outputs = model(**batch)logits = outputs.logitspredictions = torch.argmax(logits, dim=-1)metric.add_batch(predictions=predictions, references=batch["labels"])metric.compute()# 輸出：
{'accuracy': 0.8431372549019608, 'f1': 0.8907849829351535}

同樣地，結果會有隨機性，但應該差不太多。

使用accelerate庫加速訓練

之前的訓練腳本是工作在單個 CPU 或單個 GPU 上的，通過使用 accelerate 庫 ，只需少量改動就可以運行在多 GPU/TPU 上。從創建訓練/驗證數據集開始，以下是手動訓練循環的代碼：

from transformers import AdamW, AutoModelForSequenceClassification, get_schedulermodel = AutoModelForSequenceClassification.from_pretrained(checkpoint, num_labels=2)
optimizer = AdamW(model.parameters(), lr=3e-5)device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
model.to(device)num_epochs = 3
num_training_steps = num_epochs * len(train_dataloader)
lr_scheduler = get_scheduler("linear",optimizer=optimizer,num_warmup_steps=0,num_training_steps=num_training_steps,
)progress_bar = tqdm(range(num_training_steps))model.train()
for epoch in range(num_epochs):for batch in train_dataloader:batch = {k: v.to(device) for k, v in batch.items()}outputs = model(**batch)loss = outputs.lossloss.backward()optimizer.step()lr_scheduler.step()optimizer.zero_grad()progress_bar.update(1)

以下是改為 accelerate 加速所改動的部分：

+ from accelerate import Acceleratorfrom transformers import AdamW, AutoModelForSequenceClassification, get_scheduler+ accelerator = Accelerator()model = AutoModelForSequenceClassification.from_pretrained(checkpoint, num_labels=2)optimizer = AdamW(model.parameters(), lr=3e-5)- device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
- model.to(device)+ train_dataloader, eval_dataloader, model, optimizer = accelerator.prepare(
+     train_dataloader, eval_dataloader, model, optimizer
+ )num_epochs = 3num_training_steps = num_epochs * len(train_dataloader)lr_scheduler = get_scheduler("linear",optimizer=optimizer,num_warmup_steps=0,num_training_steps=num_training_steps)progress_bar = tqdm(range(num_training_steps))model.train()for epoch in range(num_epochs):for batch in train_dataloader:
-         batch = {k: v.to(device) for k, v in batch.items()}outputs = model(**batch)loss = outputs.loss
-         loss.backward()
+         accelerator.backward(loss)optimizer.step()lr_scheduler.step()optimizer.zero_grad()progress_bar.update(1)

添加的第一行是導包。第二行實例化一個 Accelerator 對象，該對象將查看環境并初始化分布式設置。Accelerate 庫會處理數據在設備上的存放，因此可以刪除將模型放置在設備上的那一行（或者，也可以將 device 改為 accelerator.device）。

然后，將數據加載器、模型和優化器送入到 accelerator.prepare()。這將把這些對象包裝在適當的容器中，以確保分布式訓練正常工作。最后一處的改動是刪除將 batch 放在設備上的那一行（同樣，如果想保留該行，可以將其更改為使用 accelerator.device 并將 loss.backward() 改為 accelerator.backward(loss) 。

以下是完成的使用 accelerate 庫加速的代碼：

from accelerate import Accelerator
from transformers import AdamW, AutoModelForSequenceClassification, get_scheduleraccelerator = Accelerator()model = AutoModelForSequenceClassification.from_pretrained(checkpoint, num_labels=2)
optimizer = AdamW(model.parameters(), lr=3e-5)train_dl, eval_dl, model, optimizer = accelerator.prepare(train_dataloader, eval_dataloader, model, optimizer
)num_epochs = 3
num_training_steps = num_epochs * len(train_dl)
lr_scheduler = get_scheduler("linear",optimizer=optimizer,num_warmup_steps=0,num_training_steps=num_training_steps,
)progress_bar = tqdm(range(num_training_steps))model.train()
for epoch in range(num_epochs):for batch in train_dl:outputs = model(**batch)loss = outputs.lossaccelerator.backward(loss)optimizer.step()lr_scheduler.step()optimizer.zero_grad()progress_bar.update(1)

將上述代碼放在 train.py 腳本中，該腳本可以在任何類型的分布式設置上運行。如果要測試自己的分布式設置，運行：

accelerate config

根據提示進行一些配置，并保存到配置文件中，然后運行：

accelerate launch train.py

這就將啟動分布式訓練。

如果想要在 Notebook （比如 Colab ）中運行，只需將上述代碼封裝到 training_function() 中，然后運行：

from accelerate import notebook_launchernotebook_launcher(training_function)

更多示例可參考：Accelerate repo.

總結

概括一下，在本章中，我們：

• 了解 Hub 中的 datasets

• 學習了如何加載和預處理數據集，包括使用動態填充和 collator

• 實現了自己對模型的微調和評估

• 實現了更底層的（基于純 Pytorch）的訓練循環

• 使用 accelerate 進行分布式訓練