| from torch import nn, Tensor |
| import torch.nn.functional as F |
| import timm |
| from typing import Union, Optional |
|
|
| from ..utils import BasicBlock, Bottleneck, make_resnet_layers |
| from ..utils import _init_weights |
|
|
|
|
| model_configs = { |
| "resnet18.tv_in1k": { |
| "decoder_channels": [512, 256, 128], |
| }, |
| "resnet34.tv_in1k": { |
| "decoder_channels": [512, 256, 128], |
| }, |
| "resnet50.tv_in1k": { |
| "decoder_channels": [512, 256, 256, 128], |
| }, |
| "resnet101.tv_in1k": { |
| "decoder_channels": [512, 512, 256, 256, 128], |
| }, |
| "resnet152.tv_in1k": { |
| "decoder_channels": [512, 512, 512, 256, 256, 128], |
| }, |
| } |
|
|
|
|
| class ResNet(nn.Module): |
| def __init__( |
| self, |
| decoder_block: Union[BasicBlock, Bottleneck], |
| backbone: str = "resnet34.tv_in1k", |
| reduction: Optional[int] = None, |
| ) -> None: |
| super().__init__() |
| assert backbone in model_configs.keys(), f"Backbone should be in {model_configs.keys()}" |
| config = model_configs[backbone] |
| encoder = timm.create_model(backbone, pretrained=True, features_only=True, out_indices=(-1,)) |
| encoder_reduction = encoder.feature_info.reduction()[-1] |
|
|
| if reduction <= 16: |
| if "resnet18" in backbone or "resnet34" in backbone: |
| encoder.layer4[0].conv1.stride = (1, 1) |
| encoder.layer4[0].downsample[0].stride = (1, 1) |
| else: |
| encoder.layer4[0].conv2.stride = (1, 1) |
| encoder.layer4[0].downsample[0].stride = (1, 1) |
| encoder_reduction = encoder_reduction // 2 |
|
|
| self.encoder = encoder |
| self.encoder_reduction = encoder_reduction |
|
|
| encoder_out_channels = self.encoder.feature_info.channels()[-1] |
|
|
| decoder_channels = config["decoder_channels"] |
| self.decoder = make_resnet_layers( |
| block=decoder_block, |
| cfg=decoder_channels, |
| in_channels=encoder_out_channels, |
| dilation=1, |
| expansion=1, |
| ) |
| self.decoder.apply(_init_weights) |
|
|
| self.reduction = self.encoder_reduction if reduction is None else reduction |
| self.channels = decoder_channels[-1] |
|
|
| def forward(self, x: Tensor) -> Tensor: |
| x = self.encoder(x)[-1] |
| if self.encoder_reduction != self.reduction: |
| x = F.interpolate(x, scale_factor=self.encoder_reduction / self.reduction, mode="bilinear") |
| x = self.decoder(x) |
|
|
| return x |
|
|
|
|
| def resnet18(reduction: int = 32) -> ResNet: |
| return ResNet(decoder_block=BasicBlock, backbone="resnet18.tv_in1k", reduction=reduction) |
|
|
|
|
| def resnet34(reduction: int = 32) -> ResNet: |
| return ResNet(decoder_block=BasicBlock, backbone="resnet34.tv_in1k", reduction=reduction) |
|
|
|
|
| def resnet50(reduction: int = 32) -> ResNet: |
| return ResNet(decoder_block=Bottleneck, backbone="resnet50.tv_in1k", reduction=reduction) |
|
|
|
|
| def resnet101(reduction: int = 32) -> ResNet: |
| return ResNet(decoder_block=Bottleneck, backbone="resnet101.tv_in1k", reduction=reduction) |
|
|
|
|
| def resnet152(reduction: int = 32) -> ResNet: |
| return ResNet(decoder_block=Bottleneck, backbone="resnet152.tv_in1k", reduction=reduction) |
|
|
