728x90
# 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_datasetsDatasetDict({
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()
728x90