การเทรนโมเดลฉบับสมบูรณ์

คราวนี้เราจะมาดูกันว่าถ้าเราอยากเขียนโค้ดให้ได้ผลลัพธ์แบบเดียวกันกับใน section ที่แล้วโดยไม่ต้องเรียกใช้ Trainer class จะต้องทำอย่างไร เราสันนิษฐานว่าคุณได้ทำกระบวนการประมวลผลข้อมูลใน section 2 มาแล้ว โค้ดข้างล่างนี้เป็นการสรุปอย่างย่อครอบคลุมถึงกระบวนการทุกอย่างที่คุณจำเป็นต้องทำ:

from datasets import load_dataset
from transformers import AutoTokenizer, DataCollatorWithPadding

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)


tokenized_datasets = raw_datasets.map(tokenize_function, batched=True)
data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

เตรียมพร้อมก่อนเทรน

ก่อนที่เราจะเริ่มเขียนลูปในการเทรนโมเดล เราจะต้องกำหนดออพเจ็กต์บางตัวก่อน โดยออพเจ็กต์ชุดแรกที่เราต้องกำหนดก็คือ dataloaders (ออพเจ็กต์สำหรับโหลดข้อมูล) ที่เราจะใช้ในการโหลดข้อมูล โดยการทำซ้ำกับหลาย ๆ batch ของข้อมูล แต่ก่อนที่คุณจะกำหนด dataloaders เหล่านี้ เราจะต้องทำกระบวนการ postprocessing บางอย่างกับ tokenized_datasets ของเราก่อน เพื่อทำกระบวนการบางอย่างที่ Trainer ได้จัดการให้เราโดยอัตโนมัติ ซึ่งกระบวนการเหล่านั้นได้แก่:

• ลบคอลัมน์ที่มีข้อมูลที่โมเดลไม่ต้องการใช้ (เช่น คอลัมน์ sentence1 และ sentence2)
• เปลี่ยนชื่อคอลัมน์ label เป็น labels (เพราะว่าโมเดลคาดหวังอากิวเมนต์ชื่อว่า labels)
• กำหนดรูปแบบของ datasets ให้ส่งผลลัพธ์ออกมาเป็น PyTorach tensors แทนที่จะเป็น lists

tokenized_datasets ของเรามีเมธอดสำหรับการจัดการแต่ละขั้นตอนดังนี้:

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"]

เมื่อเราทำขั้นตอนนี้เสร็จแล้ว เราก็สามารถกำหนด dataloaders ของเราได้อย่างง่ายดาย ดังนี้:

from torch.utils.data import DataLoader

train_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
)

เพื่อจะตรวจสอบอย่างรวดเร็วว่าไม่มีข้อผิดพลาดจากการประมวลผลข้อมูล เราสามารถลองเรียกข้อมูล batch หนึ่งมาดูได้ดังนี้:

for batch in train_dataloader:
    break
{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])}

ควรระวังไว้ว่า shape ที่คุณได้อาจจะแตกต่างไปจากนี้เล็กน้อย เนื่องจากเราได้กำหนดค่าให้ training dataloader มีการทำ shuffle=True และเราได้เติมข้อมูลให้ยาวเท่ากับข้อมูลตัวที่ยาวที่สุดใน batch

ตอนนี้เราก็ทำกระบวนการประมวลผลข้อมูลเสร็จแล้ว (สำหรับนักปฏิบัติ ML แล้วนี่เป็นเป้าหมายที่น่าพึงพอใจทีเดียว แต่ยังไม่ได้ให้ผลลัพธ์อะไรออกมาเป็นรูปธรรมนะ) ลองกลับมาดูที่โมเดลกัน เราจะสร้างโมเดลขึ้นมาแบบเดียวกับที่เราทำใน section ที่แล้ว:

from transformers import AutoModelForSequenceClassification

model = 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])

โมเดล 🤗 Transformers ทุกตัวจะให้ผลลัพธ์ค่า loss ออกมาด้วยถ้าหากเราใส่ข้อมูล labels เข้าไปด้วย และเรายังได้ผลลัพธ์ออกมาเป็น logits (ได้ออกมาเป็น 2 ค่าสำหรับแต่ละ input ใน batch ของเรา ดังนั้น logits จะเป็น tensor ที่มีขนาด 8 x 2)

เราเกือบจะพร้อมสำหรับการเขียนลูปในการเทรนแล้ว! เราแค่ต้องการอีกสองสิ่งเท่านั้นเอง: optimizer (ตัวปรับปรุงให้การเทรนราบรื่นขึ้น) และ learning rate scheduler (ตัวกำหนดค่า learning rate ตามเวลา) เนื่องจากตอนนี้เราพยายามจะเลียนแบบสิ่งที่ Trainer ทำไว้ เราก็จะใช้ค่าเริ่มต้นที่เหมือนกัน โดยตัว optimizer ที่ Trainer ใช้คือ AdamW ซึ่งเป็นตัวเดียวกันกับ Adam แต่มีการพลิกแพลงในส่วนของ weight decay regularization (ดูเพิ่มเติมที่ “Decoupled Weight Decay Regularization” โดย Ilya Loshchilov and Frank Hutter):

from transformers import AdamW

optimizer = AdamW(model.parameters(), lr=5e-5)

เราก็มาถึงขั้นตอนสุดท้าย เราจะต้องกำหนดตัว learning rate scheduler ซึ่งมีค่าเริ่มต้นให้ learningrate มีการ decay แบบเชิงเส้น โดยมีการลดค่าจากค่า learning rate ที่สูงที่สุด (5e-5) ไปเรื่อย ๆ จนมีค่าเป็น 0 เพื่อจะกำหนดค่าให้ถูกต้อง เราจะต้องรู้ว่าการเทรนครั้งนี้มีการเทรนจำนวนทั้งสิ้นกี่ step ซึ่งคำนวณได้จากจำนวน epochs ที่เราจะเทรน คูณด้วยจำนวน training batches (ซึ่งก็คือความยาวของ training dataloader ของเรา) โดยตัว Trainer มีค่าเริ่มต้นในการเทรนอยู่ที่ 3 epochs เราก็จะยึดตามค่านี้:

from transformers import get_scheduler

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,
)
print(num_training_steps)

1377

ลูปในการเทรน

ข้อควรคำนึงถึงข้อสุดท้าย: เราจะต้องการใช้ GPU ถ้าหากเรามีการติดตั้งไว้ (ถ้าเราเทรนบน CPU จะต้องใช้เวลาหลายชั่วโมงแทนที่จะเป็นเวลาไม่กี่นาที) เพื่อกำหนดให้มีการใช้ GPU ทุกครั้งที่เป็นไปได้ เราสามารถกำหนด device ที่เราจะใส่โมเดลและ batches ของข้อมูลของเราลงไปได้ดังนี้:

import torch

device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
model.to(device)
device

device(type='cuda')

ตอนนี้เราก็พร้อมจะเทรนโมเดลแล้ว! เพื่อให้เราพอคาดการณ์ได้ว่าการเทรนจะใช้เวลานานเท่าไร เราจึงเพิ่มแถบแสดงสถานะความคืบหน้าตามจำนวน step ในการเทรน โดยใช้ไลบรารี่ tqdm ดังนี้:

from tqdm.auto import tqdm

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()

        optimizer.step()
        lr_scheduler.step()
        optimizer.zero_grad()
        progress_bar.update(1)

คุณจะสามารถเห็นได้ว่าแก่นของลูปในการเทรนนั้นก็เหมือนกับตัวที่เราแสดงให้ดูในบทนำ เรายังไม่ได้กำหนดให้มีการรายงานค่าใด ๆ ออกมา ดังนั้นลูปในการเทรนตัวนี้จึงไม่ได้บอกอะไรเราเลยว่าโมเดลมีประสิทธิภาพเป็นอย่างไร เราจึงจำเป็นต้องเขียนลูปในการประเมินผลโมเดล (evaluation loop) ด้วย

ลูปในการประเมินผลโมเดล (evaluation loop)

เหมือนกับที่เราได้ทำไว้ก่อนหน้านี้ เราสามารถเรียกใช้ metric จากไลบรารี่ 🤗 Evaluate ได้เลย เราได้เห็นเมธอด metric.compute() มาแล้ว แต่ metrics ยังสามารถรวบรวมผลมาเป็น batches ให้เราได้ด้วย โดยใช้เมธอด add_batch() โดยเมื่อเรารวบรวมผลมาจากทุก batches แล้ว เราก็จะคำนวณผลลัพธ์สุดท้ายได้โดยใช้เมธอด metric.compute() โค้ดข้างล่างนี้เป็นตัวอย่างการทำทุกอย่างที่เรากล่าวมานี้ในลูปสำหรับประเมินผลโมเดล:

import evaluate

metric = 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.logits
    predictions = torch.argmax(logits, dim=-1)
    metric.add_batch(predictions=predictions, references=batch["labels"])

metric.compute()

{'accuracy': 0.8431372549019608, 'f1': 0.8907849829351535}

ผลลัพธ์ที่ได้อาจแตกต่างไปเล็กน้อยเนื่องจากมีการสุ่มค่า weight ตอนสร้าง model head และมีการสลับข้อมูลแบบสุ่ม แต่ผลที่ได้ก็ควรจะใกล้เคียงกัน

เร่งความเร็วให้ลูปในการเทรนของคุณด้วย 🤗 Accelerate

ลูปในการเทรนที่เรากำหนดขึ้นก่อนหน้านี้ทำงานได้ดีบน CPU หรือ GPU ตัวเดียว แต่การใช้ไลบรารี่ 🤗 Accelerate และเพิ่มการปรับค่าอีกเพียงเล็กน้อย จะช่วยให้เราสามารถเทรนบน distributed setup ที่มีการใช้ GPUs หรือ TPUs หลายตัวได้ โดยเริ่มต้นจากการสร้าง training และ validation dataloaders ลูปในการเทรนแบบ manual ของเรามีลักษณะดังนี้:

from transformers import AdamW, AutoModelForSequenceClassification, get_scheduler

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)

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.loss
        loss.backward()

        optimizer.step()
        lr_scheduler.step()
        optimizer.zero_grad()
        progress_bar.update(1)

และต่อไปนี้คือสิ่งที่ต้องปรับแก้:

+ from accelerate import Accelerator
  from 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 = 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.loss
-         loss.backward()
+         accelerator.backward(loss)

          optimizer.step()
          lr_scheduler.step()
          optimizer.zero_grad()
          progress_bar.update(1)

โค้ดบรรทัดแรกที่ต้องเพิ่มเข้ามาเป็นส่วนของการ import โดยบรรทัดที่สองเป็นการสร้างออพเจ็กต์ Accelerator ที่จะตรวจสอบ environment ของคุณและสร้าง distributed setup ที่เหมาะสมขึ้นมาให้ โดย 🤗 Accelerate จะช่วยจัดการ device ให้คุณ คุณจึงสามารถเอาโค้ดส่วนที่คุณใส่โมเดลเข้าไปใน device ออกได้ (หรือถ้าคุณอยากคงไว้ ก็เปลี่ยนจาก device เป็น accelerator.device)

จากนัั้นก็มีการทำงานหลัก ๆ ในบรรทัดที่ส่ง dataloaders, โมเดล และ optimizer เข้าไปที่ accelerator.prepare() ซึ่งเป็นการ wrap ออพเจ็กต์เหล่านี้ให้อยู่ใน container ที่เหมาะสม และทำให้แน่ใจว่า distributed training ของคุณทำงานได้ตามที่ตั้งใจไว้ การเปลี่ยนแปลงส่วนที่เหลือคือการเอาโค้ดส่วนที่คุณใส่ batch เข้าไปใน device ออก (และอีกครั้ง ถ้าคุณอยากคงไว้ ก็เปลี่ยนจาก device เป็น accelerator.device และแก้จาก loss.backward() เป็น accelerator.backward(loss))

ถ้าคุณอยากคัดลองและวางโค้ดเพื่อทดลองดู โค้ดข้างล่างนี้คือตัวอย่างของลูปในการเทรนโดยใช้ 🤗 Accelerate แบบสมบูรณ์:

from accelerate import Accelerator
from transformers import AdamW, AutoModelForSequenceClassification, get_scheduler

accelerator = 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.loss
        accelerator.backward(loss)

        optimizer.step()
        lr_scheduler.step()
        optimizer.zero_grad()
        progress_bar.update(1)

การใส่โค้ดข้างบนนี้เข้าไปในสคริปต์ train.py จะทำให้สคริปต์ของคุณรันได้ไม่ว่าจะมี distributed setup เป็นแบบใดก็ตาม เพื่อจะลองบน distributed setup ของคุณ ให้รันคำสั่งนี้:

accelerate config

ซึ่งจะให้คุณตอบคำถาม 2-3 ข้อ และใส่คำตอบของคุณลงไปในไฟล์ configuration ที่ใช้ในคำสั่งนี้:

accelerate launch train.py

ซึ่งจะเริ่มการเทรนโมเดลแบบ distributed

ถ้าคุณอยากลองโค้ดนี้บน Notebook (เช่น ทดลองกับ TPUs บน Colab) แค่วางโค้ดนี้ลงไปใน training_function() และรัน cell สุดท้ายด้วยโค้ดนี้:

from accelerate import notebook_launcher

notebook_launcher(training_function)

คุณสามารถศึกษาจากตัวอย่างอื่น ๆ เพิ่มเติม ได้ใน 🤗 Accelerate repo