timm integration

flash/vision/backbones.py

@@ -15,10 +15,14 @@
 
 import torchvision
 from pytorch_lightning.utilities import _BOLTS_AVAILABLE, rank_zero_warn
+from pytorch_lightning.utilities import _module_available
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from torch import nn as nn
 from torchvision.models.detection.backbone_utils import resnet_fpn_backbone
 
+if _module_available('timm'):
+    import timm
+
 if _BOLTS_AVAILABLE:
     from pl_bolts.models.self_supervised import SimCLR, SwAV
 
@@ -70,6 +74,9 @@ def backbone_and_num_features(
     if model_name in TORCHVISION_MODELS:
         return torchvision_backbone_and_num_features(model_name, pretrained)
 
+    if model_name in timm.list_models():
+        return timm_backbone_and_num_features(model_name, pretrained)
+
     raise ValueError(f"{model_name} is not supported yet.")
 
 
@@ -140,3 +147,21 @@ def torchvision_backbone_and_num_features(model_name: str, pretrained: bool = Tr
         return backbone, num_features
 
     raise ValueError(f"{model_name} is not supported yet.")
+
+
+def timm_backbone_and_num_features(model_name: str, pretrained: bool = True) -> Tuple[nn.Module, int]:
+    """
+    >>> timm_backbone_and_num_features('resnet18')  # doctest: +ELLIPSIS
+    (ResNet(...), 512)
+    >>> timm_backbone_and_num_features('mobilenetv3_large_100')  # doctest: +ELLIPSIS
+    (MobileNetV3(...), 1280)
+    """
+
+    if model_name in timm.list_models():
+        backbone = timm.create_model(model_name, pretrained=pretrained, num_classes=0,
+                                     global_pool='')
+        num_features = backbone.num_features
+        return backbone, num_features
+
+    raise ValueError(
+        f"{model_name} is not supported in timm yet. https://rwightman.github.io/pytorch-image-models/models/")

requirements.txt

@@ -17,3 +17,4 @@ sentencepiece>=0.1.95
 lightning-bolts==0.3.2  # todo: we shall align with proper release
 filelock  # comes with 3rd-party dependency
 pycocotools>=2.0.2 ; python_version >= "3.7"
+timm~=0.4.5