timm.py

from flash.core.registry import FlashRegistry
from functools import partial
import timm


TIMM_BACKBONES_REGISTRY = FlashRegistry("backbones")
for model_name in timm.list_models():
  
  def _fn_timm(model_name, pretrained = True) -> Tuple[nn.Module, int]:
        backbone = timm.create_model(
            model_name, pretrained=pretrained, num_classes=num_classes
        )
        num_features = backbone.num_features
        return backbone, num_features
  
  TIMM_BACKBONES_REGISTRY(fn=partial(_fn_timm, model_name), name=model_name)