| 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) | |
| 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): | |
| return self.model(tensor) | |