Spaces:
Sleeping
Sleeping
File size: 2,121 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 |
from timm import create_model, list_models
from torch import nn, Tensor
import torch.nn.functional as F
from typing import Optional
from warnings import warn
class TIMMEncoder(nn.Module):
def __init__(
self,
backbone: str,
reduction: Optional[int] = None,
) -> None:
super().__init__()
assert backbone in list_models(), f"Backbone {backbone} not available in timm"
encoder = create_model(backbone, pretrained=True, features_only=True, out_indices=[-1])
encoder_reduction = encoder.feature_info.reduction()[-1]
if reduction <= 16:
if "resnet" in backbone:
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
elif "mobilenetv2" in backbone:
encoder.blocks[5][0].conv_dw.stride = (1, 1)
encoder_reduction = encoder_reduction // 2
elif "densenet" in backbone:
encoder.features_transition3.pool = nn.Identity()
encoder_reduction = encoder_reduction // 2
else:
warn(f"Reduction for {backbone} not handled. Using default reduction of {encoder_reduction}")
self.encoder = encoder
self.encoder_reduction = encoder_reduction
self.reduction = self.encoder_reduction if reduction is None else reduction
self.channels = self.encoder.feature_info.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")
return x
def _timm_encoder(backbone: str, reduction: Optional[int] = None) -> TIMMEncoder:
return TIMMEncoder(backbone, reduction)
|