chunked/content_aware_chunking/_custom_models/chunk_27.txt

Then the
model is defined from the configuration by passing everything to the ResNet 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)

For the model that will classify images, we just change the forward method:

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

In both cases, notice how we inherit from PreTrainedModel and call the superclass initialization with the config
(a bit like when you write a regular torch.nn.Module).