Compartilhando modelos customizados

A biblioteca 🤗 Transformers foi projetada para ser facilmente extensível. Cada modelo é totalmente codificado em uma determinada subpasta do repositório sem abstração, para que você possa copiar facilmente um arquivo de modelagem e ajustá-lo às suas necessidades.

Se você estiver escrevendo um modelo totalmente novo, pode ser mais fácil começar do zero. Neste tutorial, mostraremos como escrever um modelo customizado e sua configuração para que possa ser usado com Transformers, e como você pode compartilhá-lo com a comunidade (com o código em que se baseia) para que qualquer pessoa possa usá-lo, mesmo se não estiver presente na biblioteca 🤗 Transformers.

Ilustraremos tudo isso em um modelo ResNet, envolvendo a classe ResNet do biblioteca timm em um PreTrainedModel.

Escrevendo uma configuração customizada

Antes de mergulharmos no modelo, vamos primeiro escrever sua configuração. A configuração de um modelo é um objeto que terá todas as informações necessárias para construir o modelo. Como veremos na próxima seção, o modelo só pode ter um config para ser inicializado, então realmente precisamos que esse objeto seja o mais completo possível.

Em nosso exemplo, pegaremos alguns argumentos da classe ResNet que podemos querer ajustar. Diferentes configurações nos dará os diferentes tipos de ResNets que são possíveis. Em seguida, apenas armazenamos esses argumentos, após verificar a validade de alguns deles.

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` must be 'basic' or bottleneck', got {block}.")
        if stem_type not in ["", "deep", "deep-tiered"]:
            raise ValueError(f"`stem_type` must be '', 'deep' or 'deep-tiered', got {block}.")

        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)

As três coisas importantes a serem lembradas ao escrever sua própria configuração são:

• você tem que herdar de PretrainedConfig,
• o __init__ do seu PretrainedConfig deve aceitar quaisquer kwargs,
• esses kwargs precisam ser passados para a superclasse __init__.

A herança é para garantir que você obtenha todas as funcionalidades da biblioteca 🤗 Transformers, enquanto as outras duas restrições vêm do fato de um PretrainedConfig ter mais campos do que os que você está configurando. Ao recarregar um config com o método from_pretrained, esses campos precisam ser aceitos pelo seu config e então enviados para a superclasse.

Definir um model_type para sua configuração (aqui model_type="resnet") não é obrigatório, a menos que você queira registrar seu modelo com as classes automáticas (veja a última seção).

Com isso feito, você pode facilmente criar e salvar sua configuração como faria com qualquer outra configuração de modelo da biblioteca. Aqui está como podemos criar uma configuração resnet50d e salvá-la:

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

Isso salvará um arquivo chamado config.json dentro da pasta custom-resnet. Você pode então recarregar sua configuração com o método from_pretrained:

resnet50d_config = ResnetConfig.from_pretrained("custom-resnet")

Você também pode usar qualquer outro método da classe PretrainedConfig, como push_to_hub() para carregar diretamente sua configuração para o Hub.

Escrevendo um modelo customizado

Agora que temos nossa configuração ResNet, podemos continuar escrevendo o modelo. Na verdade, escreveremos dois: um que extrai os recursos ocultos de um lote de imagens (como BertModel) e um que é adequado para classificação de imagem (como BertForSequenceClassification).

Como mencionamos antes, escreveremos apenas um wrapper solto do modelo para mantê-lo simples para este exemplo. A única coisa que precisamos fazer antes de escrever esta classe é um mapa entre os tipos de bloco e as classes de bloco reais. Então o modelo é definido a partir da configuração passando tudo para a classe ResNet:

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)

Para o modelo que irá classificar as imagens, vamos apenas alterar o método forward:

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.cross_entropy(logits, labels)
            return {"loss": loss, "logits": logits}
        return {"logits": logits}

Em ambos os casos, observe como herdamos de PreTrainedModel e chamamos a inicialização da superclasse com o config (um pouco parecido quando você escreve um torch.nn.Module). A linha que define o config_class não é obrigatória, a menos que você deseje registrar seu modelo com as classes automáticas (consulte a última seção).

Você pode fazer com que seu modelo retorne o que você quiser,porém retornando um dicionário como fizemos para ResnetModelForImageClassification, com a função de perda incluída quando os rótulos são passados, vai tornar seu modelo diretamente utilizável dentro da classe Trainer. Você pode usar outro formato de saída, desde que esteja planejando usar seu próprio laço de treinamento ou outra biblioteca para treinamento.

Agora que temos nossa classe do modelo, vamos criar uma:

resnet50d = ResnetModelForImageClassification(resnet50d_config)

Novamente, você pode usar qualquer um dos métodos do PreTrainedModel, como save_pretrained() ou push_to_hub(). Usaremos o segundo na próxima seção e veremos como enviar os pesos e o código do nosso modelo. Mas primeiro, vamos carregar alguns pesos pré-treinados dentro do nosso modelo.

Em seu próprio caso de uso, você provavelmente estará treinando seu modelo customizado em seus próprios dados. Para este tutorial ser rápido, usaremos a versão pré-treinada do resnet50d. Como nosso modelo é apenas um wrapper em torno dele, será fácil de transferir esses pesos:

import timm

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

Agora vamos ver como ter certeza de que quando fazemos save_pretrained() ou push_to_hub(), o código do modelo é salvo.

Enviando o código para o Hub

Primeiro, certifique-se de que seu modelo esteja totalmente definido em um arquivo .py. Ele pode contar com importações relativas para alguns outros arquivos desde que todos os arquivos estejam no mesmo diretório (ainda não suportamos submódulos para este recurso). Para o nosso exemplo, vamos definir um arquivo modeling_resnet.py e um arquivo configuration_resnet.py em uma pasta no diretório de trabalho atual chamado resnet_model. O arquivo de configuração contém o código para ResnetConfig e o arquivo de modelagem contém o código do ResnetModel e ResnetModelForImageClassification.

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

O __init__.py pode estar vazio, apenas está lá para que o Python detecte que o resnet_model possa ser usado como um módulo.

Observe que você pode reutilizar (ou subclasse) uma configuração/modelo existente.

Para compartilhar seu modelo com a comunidade, siga estas etapas: primeiro importe o modelo ResNet e a configuração do arquivos criados:

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

Então você tem que dizer à biblioteca que deseja copiar os arquivos de código desses objetos ao usar o save_pretrained e registrá-los corretamente com uma determinada classe automáticas (especialmente para modelos), basta executar:

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

Observe que não há necessidade de especificar uma classe automática para a configuração (há apenas uma classe automática, AutoConfig), mas é diferente para os modelos. Seu modelo customizado pode ser adequado para muitas tarefas diferentes, então você tem que especificar qual das classes automáticas é a correta para o seu modelo.

Em seguida, vamos criar a configuração e os modelos como fizemos antes:

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())

Agora para enviar o modelo para o Hub, certifique-se de estar logado. Ou execute no seu terminal:

huggingface-cli login

ou a partir do notebook:

from huggingface_hub import notebook_login

notebook_login()

Você pode então enviar para seu próprio namespace (ou uma organização da qual você é membro) assim:

resnet50d.push_to_hub("custom-resnet50d")

Além dos pesos do modelo e da configuração no formato json, isso também copiou o modelo e configuração .py na pasta custom-resnet50d e carregou o resultado para o Hub. Você pode conferir o resultado neste repositório de modelos.

Consulte o tutorial de compartilhamento para obter mais informações sobre o método push_to_hub.

Usando um modelo com código customizado

Você pode usar qualquer configuração, modelo ou tokenizador com arquivos de código customizados em seu repositório com as classes automáticas e o método from_pretrained. Todos os arquivos e códigos carregados no Hub são verificados quanto a malware (consulte a documentação de Segurança do Hub para obter mais informações), mas você ainda deve revisar o código do modelo e o autor para evitar a execução de código malicioso em sua máquina. Defina trust_remote_code=True para usar um modelo com código customizado:

from transformers import AutoModelForImageClassification

model = AutoModelForImageClassification.from_pretrained("sgugger/custom-resnet50d", trust_remote_code=True)

Também é fortemente recomendado passar um hash de confirmação como uma revisão para garantir que o autor dos modelos não atualize o código com novas linhas maliciosas (a menos que você confie totalmente nos autores dos modelos).

commit_hash = "ed94a7c6247d8aedce4647f00f20de6875b5b292"
model = AutoModelForImageClassification.from_pretrained(
    "sgugger/custom-resnet50d", trust_remote_code=True, revision=commit_hash
)

Observe que ao navegar no histórico de commits do repositório do modelo no Hub, há um botão para copiar facilmente o commit hash de qualquer commit.

Registrando um modelo com código customizado para as classes automáticas

Se você estiver escrevendo uma biblioteca que estende 🤗 Transformers, talvez queira estender as classes automáticas para incluir seus próprios modelos. Isso é diferente de enviar o código para o Hub no sentido de que os usuários precisarão importar sua biblioteca para obter os modelos customizados (ao contrário de baixar automaticamente o código do modelo do Hub).

Desde que sua configuração tenha um atributo model_type diferente dos tipos de modelo existentes e que as classes do seu modelo tenha os atributos config_class corretos, você pode simplesmente adicioná-los às classes automáticas assim:

from transformers import AutoConfig, AutoModel, AutoModelForImageClassification

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

Observe que o primeiro argumento usado ao registrar sua configuração customizada para AutoConfig precisa corresponder ao model_type de sua configuração customizada. E o primeiro argumento usado ao registrar seus modelos customizados, para qualquer necessidade de classe de modelo automático deve corresponder ao config_class desses modelos.