| |
| |
| |
| |
| import torch |
| import timm |
| import logging |
|
|
| from typing import TYPE_CHECKING, TypeVar, Generic, Optional, Mapping, Any |
| from typing_extensions import override |
| from torch import Tensor |
| from transformers import PreTrainedModel, PretrainedConfig |
| from timm.models.resnet import BasicBlock, Bottleneck, ResNet |
|
|
| from .configuration import ResnetConfig |
|
|
| if TYPE_CHECKING: |
| ... |
|
|
| logger = logging.getLogger(__name__) |
| BLOCK_MAPPING = {"basic": BasicBlock, "bottleneck": Bottleneck} |
|
|
| ConfigClass = TypeVar("ConfigClass", bound=PretrainedConfig) |
|
|
|
|
| class ResnetModel(PreTrainedModel, Generic[ConfigClass]): |
| """The custom resnet class.""" |
|
|
| config_class: ConfigClass = ResnetConfig |
| """""" |
|
|
| def __init__(self, config: ConfigClass): |
| super().__init__(config=config) |
|
|
| |
| self.config: ResnetConfig = config |
|
|
| block_layer = BLOCK_MAPPING[self.config.block_type] |
|
|
| self.model = ResNet( |
| block=block_layer, |
| layers=self.config.layers, |
| num_classes=self.config.num_classes, |
| in_chans=self.config.input_channels, |
| cardinality=self.config.cardinality, |
| base_width=self.config.base_width, |
| stem_width=self.config.stem_width, |
| stem_type=self.config.stem_type, |
| avg_down=self.config.avg_down, |
| ) |
|
|
| def load_state_dict_from_source(self, source: Optional[str] = None): |
| source = source or self.config.pretrained_state_source |
|
|
| if source == "timm": |
| pretrained_model = timm.create_model("resnet50d", pretrained=True) |
| self.load_state_dict( |
| state_dict=pretrained_model.state_dict(), |
| assign=False, |
| strict=False |
| ) |
|
|
| del pretrained_model |
| else: |
| logger.warning(f"Now, only support loading pretrained weights from 'timm', but got {source!r}.") |
|
|
| @override |
| def load_state_dict( |
| self, state_dict: Mapping[str, Any], |
| strict: bool = True, |
| assign: bool = False |
| ): |
| self.model.load_state_dict(state_dict, strict=strict, assign=assign) |
|
|
| def forward(self, tensor: Tensor) -> Tensor: |
| return self.model.forward(tensor) |
|
|
|
|
| class ResnetModelForImageClassification(ResnetModel): |
| 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, |
| ) |
|
|
| @override |
| def forward(self, tensor: Tensor, labels=None) -> dict[str, Tensor]: |
| 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} |
|
|
