IMDb 데이터셋의 label 열이 DistilBERT 모델 forward() 함수의 labels 인자로 전달되는 과정

1. 개요

허깅 페이스의 Transformers 라이브러리를 사용하여 모델을 훈련할 때 데이터셋의 label 항목이 어떤 과정을 거쳐 모델의 forward(..., labels, ...) 메소드로 전달되는지 설명합니다.

IMDb 데이터셋에서 추출한 한 개의 데이터 예시

{
    "text": "I love sci-fi...",
    "label": 0
}

DistilBertForSequenceClassification 클래스의 forward() 메소드

transformers/models/distilbert/modeling_distilbert.py

def forward(
    ...,
    labels: Optional[torch.LongTensor] = None,
    ...
)-> Union[SequenceClassifierOutput, Tuple[torch.Tensor, ...]]:
    ...

이 문서에서는 IMDb 데이터셋과 DistilBertForSequenceClassification 모델을 사용하여 설명하지만 특정 데이터셋과 모델에만 해당하는 것은 아닙니다.

이 문서에서 등장하는 주요 API는 다음과 같습니다.

datasets.default_data_collator
class datasets.Dataset
class datasets.DatasetDict
class transformers.DataCollatorWithPadding
class transformers.DistilBertForSequenceClassification
class transformers.Trainer

2. 트레이너 (Trainer)

2.1. `Trainer.init()`

데이터 콜레이터를 인자로 받아들이는 Trainer 객체 생성

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_imdb["train"],
    eval_dataset=tokenized_imdb["test"],
    processing_class=tokenizer,
    data_collator=data_collator,
    compute_metrics=compute_metrics,
)

Trainer 클래스의 생성자에서 사용자 지정 데이터 콜레이터를 쓸 것인지, 아니면 기본 데이터 콜레이터를 쓸 것인지 결정

transformers/trainer.py

class Trainer:
    def __init__(self, ..., data_collator, ...):
        ...
        default_collator = (
            DataCollatorWithPadding(processing_class)
            if processing_class is not None
            and isinstance(processing_class, (PreTrainedTokenizerBase, SequenceFeatureExtractor))
            else default_data_collator
        )
        self.data_collator = data_collator if data_collator is not None else default_collator
        ...

2.2. `Trainer.train()`

훈련 과정에서 배치 데이터를 얻기 위하여 __init__() 메소드 내에서 정했던 데이터 콜레이터를 이용

transformers/trainer.py

class Trainer:
    ...
    def train(self, ...):
        ...
        find_executable_batch_size(self._inner_training_loop, ...)
        ...
    
    def _inner_training_loop(self, batch_size=None, ...):
        ...
        train_dataloader = self.get_train_dataloader()
        ...
        for epoch in range(epochs_trained, num_train_epochs):
            epoch_dataloader = train_dataloader
            ...
            epoch_iterator = iter(epoch_dataloader)
            ...
            for _ in range(total_updates):
                batch_samples, num_items_in_batch = self.get_batch_samples(epoch_iterator, num_batches, args.device)
                ...
                for i, inputs in enumerate(batch_samples):
                    ...
                    tr_loss_step = self.training_step(model, inputs, num_items_in_batch)
                    ...
        ...
    
    def get_train_dataloader(self) -> DataLoader:
        ...
        data_collator = self.data_collator
        ...
        
    def training_step(self, model, inputs, ...) -> torch.Tensor:
        ...
        loss = self.compute_loss(model, inputs, num_items_in_batch=num_items_in_batch)
        ...

    def compute_loss(self, model, inputs, ...)
        ...
        outputs = model(**inputs)
        ...
        return (loss, outputs) if return_outputs else loss

3. 데이터 콜레이터 (Data Collator)

데이터 콜레이터는 데이터셋으로부터 배치 크기의 데이터를 추출하여 반환하는 역할을 수행합니다. 트레이너는 훈련 과정에서 데이터를 공급받기 위하여 응용 프로그램에서 직접 생성하여 지정한 데이터 콜레이터를 사용하거나, 그렇지 않으면 기본 데이터 콜레이터를 사용합니다.

Transformers 라이브러리는 다음 세 종류의 기본 데이터 콜레이터를 구현하고 있습니다.

PyTorch - torch_default_data_collator
TensorFlow - tf_default_data_collator
NumPy - numpy_default_data_collator

3.1. 데이터 콜레이터를 직접 생성하여 지정하는 경우

DataCollatorWithPadding 객체를 생성하여 Trainer 객체 생성 시 인자로 전달

data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

trainer = Trainer(
    ...
    data_collator=data_collator,
    ...
)

trainer.train()

DataCollatorWithPadding 객체 호출 시 label 항목을 labels로 변경

transformers/data/data_collator.py

class DataCollatorWithPadding:
    def __call__(self, ...):
        batch = pad_without_fast_tokenizer_warning(...)
        if "label" in batch:
            batch["labels"] = batch["label"]
	        del batch["label"]
        if "label_ids" in batch:
            batch["labels"] = batch["label_ids"]
            del batch["label_ids"]
	    return batch

3.2. 데이터 콜레이터를 지정하지 않는 경우

Trainer 객체 생성 시 data_collator 파라미터 지정하지 않음

trainer = Trainer(
    ...
    data_collator=None,
    ...
)

PyTorch의 경우 torch_default_data_collator() 함수 호출 시 label 항목을 labels로 변경

def torch_default_data_collator():
    ...
    if "label" in ...:
	    ...
        batch["labels"] = torch.tensor([f["label"] for f in features], dtype=dtype)

4. 정리

데이터 콜레이터가 데이터셋의 label 항목을 labels 항목으로 변경하고 트레이너가 labels 항목을 모델의 forward(..., labels, ...) 메소드 인자로 전달합니다.

Written with StackEdit.

천천히, 제대로

이 블로그 검색