|
from typing import Dict |
|
|
|
import timm |
|
from timm.models.resnet import BasicBlock, Bottleneck, ResNet |
|
from torch import Tensor, nn |
|
from transformers import PreTrainedModel |
|
|
|
from .configuration_resnet import ResnetConfig |
|
|
|
BLOCK_MAPPING = {"basic": BasicBlock, "bottleneck": Bottleneck} |
|
|
|
|
|
class ResnetModel(PreTrainedModel): |
|
""" |
|
The line that sets the config_class is not mandatory, |
|
unless you want to register your model with the auto classes |
|
""" |
|
|
|
config_class = ResnetConfig |
|
|
|
def __init__(self, config: ResnetConfig): |
|
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: Tensor) -> Tensor: |
|
return self.model.forward_features(tensor) |
|
|
|
|
|
class ResnetModelForImageClassification(PreTrainedModel): |
|
""" |
|
The line that sets the config_class is not mandatory, |
|
unless you want to register your model with the auto classes |
|
""" |
|
|
|
config_class = ResnetConfig |
|
|
|
def __init__(self, config: ResnetConfig): |
|
super().__init__(config) |
|
self.model = ResnetModel(config) |
|
|
|
""" |
|
You can have your model return anything you want, |
|
but returning a dictionary like we did for ResnetModelForImageClassification, |
|
with the loss included when labels are passed, |
|
will make your model directly usable inside the Trainer class. |
|
Using another output format is fine as long as you are planning on |
|
using your own training loop or another library for training. |
|
""" |
|
|
|
def forward(self, tensor: Tensor, labels=None) -> Dict[str, Tensor]: |
|
logits = self.model(tensor) |
|
if labels is not None: |
|
loss = nn.cross_entropy(logits, labels) |
|
return {"loss": loss, "logits": logits} |
|
return {"logits": logits} |
|
|
|
|
|
if __name__ == "__main__": |
|
resnet50d_config = ResnetConfig.from_pretrained("custom-resnet") |
|
resnet50d = ResnetModelForImageClassification(resnet50d_config) |
|
|
|
|
|
pretrained_model: nn.Module = timm.create_model("resnet50d", pretrained=True) |
|
resnet50d.model.load_state_dict(pretrained_model.state_dict()) |
|
|
