Spaces:
Sleeping
Sleeping
File size: 3,189 Bytes
570db9a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 |
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)
|