모델 커스터마이징

모델 커스터마이징 (Customizing models)

Transformers 모델은 커스터마이징할 수 있게 설계됐어요. ResNet 모델을 커스터마이징하고, AutoClass 지원을 켜고, Hub에 공유하는 과정을 살펴볼게요.

출처: 문서

본문

Transformers 모델은 커스터마이징할 수 있게 설계돼요. 모델 코드는 Transformers 저장소의 model 하위 폴더에 완전히 들어 있어요. 각 폴더에는 modeling.py와 configuration.py 파일이 있어요. 이 파일들을 복사해 모델 커스터마이징을 시작해요.

[!TIP] 완전히 새로운 모델을 만든다면 처음부터 시작하는 게 더 쉬울 수 있어요. 하지만 Transformers의 기존 모델과 매우 비슷한 모델이라면 같은 configuration과 model 클래스를 재사용하거나 서브클래싱하는 게 더 빨라요.

이 가이드는 ResNet 모델을 커스터마이징하고, AutoClass 지원을 켜고, Hub에 공유하는 방법을 보여줘요.

Configuration

기본 PreTrainedConfig 클래스로 주어지는 configuration은 모델을 구축하는 데 필요한 모든 정보를 담아요. 이곳에서 커스텀 ResNet 모델의 속성을 구성해요. 서로 다른 속성은 서로 다른 ResNet 모델 타입을 만들어요.

configuration을 커스터마이징하기 위한 주요 규칙은:

1. 커스텀 configuration은 PreTrainedConfig를 서브클래싱해야 해요. 그래야 커스텀 모델이 from_pretrained(), save_pretrained(), push_to_hub() 같은 Transformers 모델의 모든 기능을 갖게 돼요.
2. PreTrainedConfig의 __init__은 모든 kwargs를 받아야 하고, 그 kwargs를 슈퍼클래스의 __init__으로 넘겨야 해요. PreTrainedConfig는 커스텀 configuration에 설정된 것보다 더 많은 필드를 가지므로, from_pretrained()로 configuration을 불러올 때 그 필드들을 커스텀 configuration이 받아 슈퍼클래스로 넘겨야 해요.

[!TIP] 파라미터 중 일부의 유효성을 검사하는 게 유용해요. 아래 예시에서는 block_type과 stem_type이 미리 정의된 값 중 하나인지 확인하는 검사가 구현돼 있어요.

AutoClass 지원을 켜려면 configuration 클래스에 model_type을 추가하세요.

from transformers import PreTrainedConfig
from typing import List

class ResnetConfig(PreTrainedConfig):
    model_type = "resnet"

    def __init__(
        self,
        block_type="bottleneck",
        layers: list[int] = [3, 4, 6, 3],
        num_classes: int = 1000,
        input_channels: int = 3,
        cardinality: int = 1,
        base_width: int = 64,
        stem_width: int = 64,
        stem_type: str = "",
        avg_down: bool = False,
        **kwargs,
    ):
        if block_type not in ["basic", "bottleneck"]:
            raise ValueError(f"`block_type` must be 'basic' or bottleneck', got {block_type}.")
        if stem_type not in ["", "deep", "deep-tiered"]:
            raise ValueError(f"`stem_type` must be '', 'deep' or 'deep-tiered', got {stem_type}.")

        self.block_type = block_type
        self.layers = layers
        self.num_classes = num_classes
        self.input_channels = input_channels
        self.cardinality = cardinality
        self.base_width = base_width
        self.stem_width = stem_width
        self.stem_type = stem_type
        self.avg_down = avg_down
        super().__init__(**kwargs)

configuration을 save_pretrained()로 커스텀 모델 폴더 custom-resnet의 JSON 파일에 저장해요.

resnet50d_config = ResnetConfig(block_type="bottleneck", stem_width=32, stem_type="deep", avg_down=True)
resnet50d_config.save_pretrained("custom-resnet")

Model

커스텀 ResNet configuration으로 이제 모델을 만들고 커스터마이징할 수 있어요. 모델은 기본 PreTrainedModel 클래스를 서브클래싱해요. PreTrainedConfig처럼 PreTrainedModel에서 상속받고 슈퍼클래스를 configuration으로 초기화하면, 저장·로딩 같은 Transformers 기능이 커스텀 모델로 확장돼요.

Transformers 모델은 __init__ 메서드에서 config 객체를 받는 컨벤션을 따라요. 이는 config 객체를 서브레이어에 각각 넘기는 여러 인자로 쪼개는 대신, 전체 config를 모델 서브레이어에 전달해요.

이런 식으로 모델을 작성하면 하이퍼파라미터의 명확한 단일 출처(single source of truth)를 갖는 더 간단한 코드가 돼요. 다른 Transformers 모델의 코드를 재사용하기도 더 쉬워요.

hidden states를 출력하는 미니멀 ResNet 모델과 이미지 분류 헤드가 있는 ResNet 모델, 두 개의 ResNet 모델을 만들 거예요.

블록 타입과 클래스 사이의 매핑을 정의해요. 나머지는 configuration 클래스를 ResNet 모델 클래스에 넘겨서 만들어요.

[!TIP] AutoClass 지원을 켜려면 모델 클래스에 config_class를 추가하세요.

from transformers import PreTrainedModel
from timm.models.resnet import BasicBlock, Bottleneck, ResNet
from .configuration_resnet import ResnetConfig

BLOCK_MAPPING = {"basic": BasicBlock, "bottleneck": Bottleneck}

class ResnetModel(PreTrainedModel):
    config_class = ResnetConfig

    def __init__(self, config):
        super().__init__(config)
        block_layer = BLOCK_MAPPING[config.block_type]
        self.model = ResNet(
            block_layer,
            config.layers,
            num_classes=config.num_classes,
            in_chans=config.input_channels,
            cardinality=config.cardinality,
            base_width=config.base_width,
            stem_width=config.stem_width,
            stem_type=config.stem_type,
            avg_down=config.avg_down,
        )

    def forward(self, tensor):
        return self.model.forward_features(tensor)

forward 메서드는 labels가 있을 때 각 logit에 대한 손실을 계산하도록 다시 써야 해요. 그 외에는 ResNet 모델 클래스가 동일해요.

[!TIP] AutoClass 지원을 켜려면 모델 클래스에 config_class를 추가하세요.

import torch

class ResnetModelForImageClassification(PreTrainedModel):
    config_class = ResnetConfig

    def __init__(self, config):
        super().__init__(config)
        block_layer = BLOCK_MAPPING[config.block_type]
        self.model = ResNet(
            block_layer,
            config.layers,
            num_classes=config.num_classes,
            in_chans=config.input_channels,
            cardinality=config.cardinality,
            base_width=config.base_width,
            stem_width=config.stem_width,
            stem_type=config.stem_type,
            avg_down=config.avg_down,
        )

    def forward(self, tensor, labels=None):
        logits = self.model(tensor)
        if labels is not None:
            loss = torch.nn.functional.cross_entropy(logits, labels)
            return {"loss": loss, "logits": logits}
        return {"logits": logits}

모델은 어떤 출력 포맷이든 반환할 수 있어요. labels가 있을 때 손실을 포함한 딕셔너리(ResnetModelForImageClassification처럼)를 반환하면 커스텀 모델이 Trainer와 호환돼요. 다른 출력 포맷에서는 자체 학습 루프나 다른 라이브러리가 필요해요.

configuration으로 커스텀 모델 클래스를 인스턴스화해요.

resnet50d = ResnetModelForImageClassification(resnet50d_config)

이 시점에서 모델에 사전학습된 가중치를 불러오거나 처음부터 학습시킬 수 있어요. 이 가이드에서는 사전학습된 가중치를 불러올 거예요.

timm 라이브러리에서 사전학습된 가중치를 불러온 뒤, load_state_dict로 그 가중치를 커스텀 모델로 옮겨요.

import timm

pretrained_model = timm.create_model("resnet50d", pretrained=True)
resnet50d.model.load_state_dict(pretrained_model.state_dict())

AutoClass

AutoClass API는 주어진 모델에 맞는 올바른 아키텍처를 자동으로 불러오기 위한 단축키예요. 커스텀 모델을 불러오는 사용자를 위해 이를 활성화하는 게 편리해요.

configuration 클래스에 model_type 속성(기존 모델 타입과 달라야 함)이 있고, 모델 클래스에 config_class 속성이 있는지 확인해요. register() 메서드로 커스텀 configuration과 모델을 AutoClass API에 추가해요.

[!TIP] AutoConfig.register()의 첫 번째 인자는 커스텀 configuration 클래스의 model_type 속성과 일치해야 하고, AutoModel.register()의 첫 번째 인자는 커스텀 모델 클래스의 config_class와 일치해야 해요.

from transformers import AutoConfig, AutoModel, AutoModelForImageClassification

AutoConfig.register("resnet", ResnetConfig)
AutoModel.register(ResnetConfig, ResnetModel)
AutoModelForImageClassification.register(ResnetConfig, ResnetModelForImageClassification)

커스텀 모델 코드가 이제 AutoClass API와 호환돼요. 사용자들은 AutoModel 또는 AutoModelForImageClassification 클래스로 모델을 불러올 수 있어요.

업로드 (Upload)

커스텀 모델을 Hub에 업로드하면 다른 사용자들이 쉽게 불러와 사용할 수 있어요.

모델 디렉터리가 아래처럼 올바르게 구성되어 있는지 확인해요. 디렉터리는 다음을 담아야 해요:

• modeling.py: ResnetModel과 ResnetModelForImageClassification의 코드를 담아요. 이 파일은 같은 디렉터리 안에 있는 한 상대 임포트에 의존할 수 있어요.

[!WARNING] Transformers 모델 파일을 복사할 때는 modeling.py 파일 상단의 모든 상대 임포트를 Transformers에서 임포트하도록 교체하세요.

• configuration.py: ResnetConfig의 코드를 담아요.
• __init__.py: 비어 있어도 되며, 이 파일 덕분에 Python resnet_model을 모듈로 사용할 수 있어요.

.
└── resnet_model
    ├── __init__.py
    ├── configuration_resnet.py
    └── modeling_resnet.py

모델을 공유하려면 ResNet 모델과 configuration을 임포트해요.

from resnet_model.configuration_resnet import ResnetConfig
from resnet_model.modeling_resnet import ResnetModel, ResnetModelForImageClassification

모델과 configuration 파일에서 코드를 복사해요. AutoClass 객체가 save_pretrained()로 저장되도록 register_for_auto_class() 메서드를 호출해요. 이는 configuration JSON 파일에 AutoClass 객체와 매핑을 포함하도록 수정해요.

모델의 경우 작업에 맞는 적절한 AutoModelFor 클래스를 고르세요.

ResnetConfig.register_for_auto_class()
ResnetModel.register_for_auto_class("AutoModel")
ResnetModelForImageClassification.register_for_auto_class("AutoModelForImageClassification")

모델에 작업을 두 개 이상 매핑하려면 configuration JSON 파일의 auto_map을 직접 편집해요.

"auto_map": {
    "AutoConfig": "<your-repo-name>--<config-name>",
    "AutoModel": "<your-repo-name>--<config-name>",
    "AutoModelFor<Task>": "<your-repo-name>--<config-name>",    
},

configuration과 모델을 만들고 사전학습된 가중치를 불러와요.

resnet50d_config = ResnetConfig(block_type="bottleneck", stem_width=32, stem_type="deep", avg_down=True)
resnet50d = ResnetModelForImageClassification(resnet50d_config)

pretrained_model = timm.create_model("resnet50d", pretrained=True)
resnet50d.model.load_state_dict(pretrained_model.state_dict())

이제 모델을 Hub에 푸시할 준비가 됐어요. 명령줄이나 노트북에서 Hugging Face 계정에 로그인해요.

hf auth login

from huggingface_hub import notebook_login

notebook_login()

모델에 push_to_hub()를 호출해 Hub에 업로드해요.

resnet50d.push_to_hub("custom-resnet50d")

사전학습된 가중치, configuration, modeling.py, configuration.py 파일이 모두 당신의 네임스페이스 아래 저장소에 업로드됐을 거예요.

커스텀 모델은 Transformers 모델과 같은 모델링 코드를 사용하지 않으므로, 불러오려면 from_pretrained()에 trust_remote_code=True를 추가해야 해요. 커스텀 모델 로딩에 대한 자세한 내용은 custom models 섹션을 참고하세요.

더 알아보기 (Learn more)

• 토크나이저 커스터마이징 문서
• 🤗 Transformers에 기여하기 문서