HugginFace Trainer이란

# Trainer 이란?

 

trainer는 모델을 학습시키기 위한 클래스입니다. training, fine-tuning, evaluation 모두 trainer class를 이용하여 할 수 있습니다.

trainer을 사용할 경우, TrainingArguments 를 통해 Huggingface 프레임워크에서 제공하는 기능들을 통합적으로 커스터마이징하여 모델을 손쉽게 학습시킬 수 있다는 장점이 있습니다.

trainer API를 사용하기 위해선 TrainingArguments인스턴스를 생성해야 하며, 이 때 학습에 필요한 여러 arguments들이 정의됩니다.

 

 

# 데이터셋 로드

from datasets import load_dataset
from transformers import AutoTokenizer, TrainingArguments, Trainer, AutoModelForSequenceClassification

raw_datasets = load_dataset("glue", "cola")
checkpoint = "bert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(checkpoint)

raw_datasets

DatasetDict({
    train: Dataset({
        features: ['sentence', 'label', 'idx'],
        num_rows: 8551
    })
    validation: Dataset({
        features: ['sentence', 'label', 'idx'],
        num_rows: 1043
    })
    test: Dataset({
        features: ['sentence', 'label', 'idx'],
        num_rows: 1063
    })
})

# 모델 & 토크나이저 로드

model_name_or_path = "bert-base-uncased"
model = AutoModelForSequenceClassification.from_pretrained(model_name_or_path, num_labels=2)    # COLA dataset의 라벨은 0(unacceptable)과 1(accpetable) 두 가지로 구분됨
tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)

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

tokenized_datasets = raw_datasets.map(tokenize_function, batched=True)

# Trainer 

training_args = TrainingArguments(
    output_dir='./results',              # output이 저장될 경로
    num_train_epochs=1,              # train 시킬 총 epochs
    per_device_train_batch_size=16,  # 각 device 당 batch size
    per_device_eval_batch_size=64,   # evaluation 시에 batch size
    warmup_steps=500,                # learning rate scheduler에 따른 warmup_step 설정
    weight_decay=0.01,                 # weight decay
    logging_dir='./logs',                 # log가 저장될 경로
    do_train=True,                        # train 수행여부
    do_eval=True,                        # eval 수행여부
    eval_steps=1000,
    group_by_length=False,
)

trainer = Trainer(
    model,                                                                    # 학습시킬 model
    args=training_args,                                                # TrainingArguments을 통해 설정한 arguments
    train_dataset=tokenized_datasets["train"],         # training dataset
    eval_dataset=tokenized_datasets["validation"], # validation dataset
    tokenizer=tokenizer,
)

# 모델 학습
trainer.train()