BertForSequenceClassificationのファインチューニング

What Is Fine-Tuning?

BertForSequenceClassification Overview

from transformers import BertForSequenceClassification
import torch

# 2 classes: negative (0) and positive (1)
model = BertForSequenceClassification.from_pretrained(
    'bert-base-uncased',
    num_labels=2
)
print(model.config.num_labels)   # 2
print(model.classifier)          # Linear(in=768, out=2)

Loading the IMDB Dataset

from datasets import load_dataset
from transformers import BertTokenizer

dataset = load_dataset('imdb')
print(dataset)  # DatasetDict with train/test splits

tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')

def tokenize_fn(examples):
    return tokenizer(
        examples['text'],
        truncation=True,
        padding='max_length',
        max_length=256
    )

tokenized = dataset.map(tokenize_fn, batched=True)
tokenized.set_format('torch', columns=['input_ids', 'attention_mask', 'label'])

Creating PyTorch DataLoaders

from torch.utils.data import DataLoader

train_dataset = tokenized['train'].select(range(2000))  # small subset for demo
test_dataset = tokenized['test'].select(range(500))

train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)

print('Train batches:', len(train_loader))
print('Test batches:', len(test_loader))

Setting Up the Optimizer

from torch.optim import AdamW
from transformers import get_linear_schedule_with_warmup

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

num_epochs = 2
num_steps = len(train_loader) * num_epochs
warmup_steps = int(0.1 * num_steps)

scheduler = get_linear_schedule_with_warmup(
    optimizer,
    num_warmup_steps=warmup_steps,
    num_training_steps=num_steps
)

The Fine-Tuning Training Loop

import torch

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

for epoch in range(2):
    model.train()
    total_loss = 0
    for batch in train_loader:
        optimizer.zero_grad()
        input_ids = batch['input_ids'].to(device)
        attention_mask = batch['attention_mask'].to(device)
        labels = batch['label'].to(device)

        outputs = model(input_ids=input_ids,
                        attention_mask=attention_mask,
                        labels=labels)
        loss = outputs.loss
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        scheduler.step()
        total_loss += loss.item()

    print(f'Epoch {epoch+1} loss: {total_loss/len(train_loader):.4f}')

Evaluating on the Validation Set

import torch
from sklearn.metrics import accuracy_score

model.eval()
all_preds, all_labels = [], []

with torch.no_grad():
    for batch in test_loader:
        input_ids = batch['input_ids'].to(device)
        attention_mask = batch['attention_mask'].to(device)
        labels = batch['label'].to(device)

        outputs = model(input_ids=input_ids,
                        attention_mask=attention_mask)
        preds = torch.argmax(outputs.logits, dim=1)
        all_preds.extend(preds.cpu().numpy())
        all_labels.extend(labels.cpu().numpy())

print('Accuracy:', accuracy_score(all_labels, all_preds))

Saving the Fine-Tuned Model

import os

save_dir = './bert_imdb_finetuned'
os.makedirs(save_dir, exist_ok=True)

model.save_pretrained(save_dir)
tokenizer.save_pretrained(save_dir)
print('Model saved to', save_dir)

# Reload in a new session:
# from transformers import BertForSequenceClassification, BertTokenizer
# model = BertForSequenceClassification.from_pretrained(save_dir)
# tokenizer = BertTokenizer.from_pretrained(save_dir)

Using the Trainer API

from transformers import Trainer, TrainingArguments
import numpy as np
from sklearn.metrics import accuracy_score

def compute_metrics(eval_pred):
    logits, labels = eval_pred
    preds = np.argmax(logits, axis=1)
    return {'accuracy': accuracy_score(labels, preds)}

training_args = TrainingArguments(
    output_dir='./bert_trainer',
    num_train_epochs=2,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=32,
    learning_rate=2e-5,
    evaluation_strategy='epoch',
    save_strategy='epoch',
    load_best_model_at_end=True
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=test_dataset,
    compute_metrics=compute_metrics
)
trainer.train()

Monitoring Training with TensorBoard

from transformers import TrainingArguments

training_args = TrainingArguments(
    output_dir='./bert_trainer',
    num_train_epochs=3,
    per_device_train_batch_size=16,
    learning_rate=2e-5,
    evaluation_strategy='steps',
    eval_steps=100,
    logging_steps=50,
    report_to='tensorboard',   # enable TensorBoard logging
    logging_dir='./bert_trainer/runs'
)
# Then: tensorboard --logdir ./bert_trainer/runs

Choosing How Many Layers to Freeze

from transformers import BertForSequenceClassification

model = BertForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=2)

# Freeze embedding layer and first 6 encoder layers
for name, param in model.named_parameters():
    if 'embeddings' in name or 'encoder.layer.0' in name or \
       'encoder.layer.1' in name or 'encoder.layer.2' in name:
        param.requires_grad = False

# Count trainable parameters
trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f'Trainable parameters: {trainable:,}')

Quick Check

Lesson Recap