[C] Update for timm registrys

chameleon/modules/backbones/__init__.py

@@ -1,14 +1,4 @@
-from functools import partial
-
-import timm
-
-from ...registry import BACKBONES
 from .gpunet import GPUNet
-
-timm_models = timm.list_models()
-for name in timm_models:
-    create_func = partial(timm.create_model, model_name=name)
-    BACKBONES.register_module(f'timm_{name}', module=create_func)
-
+from .timm import *
 
 __all__ = ['GPUNet']

chameleon/modules/backbones/timm.py

@@ -0,0 +1,15 @@
+import timm
+import torch.nn as nn
+
+from ...registry import BACKBONES
+
+
+class Timm:
+    @staticmethod
+    def build_model(*args, **kwargs) -> nn.Module:
+        return timm.create_model(*args, **kwargs)
+
+
+timm_models = timm.list_models()
+for name in timm_models:
+    BACKBONES.register_module(f'timm_{name}', module=Timm.build_model)

tests/registry/test_root.py

@@ -1,4 +1,5 @@
 import torch
+import torch.nn as nn
 from torchmetrics.metric import Metric
 
 from chameleon import ASPP, FPN, AdamW, AWingLoss, Conv2dBlock, GPUNet

tests/tools/test_calflops.py

@@ -5,7 +5,8 @@
 
 
 def test_calcualte_flops():
-    model = BACKBONES.build({'name': 'timm_resnet50'})
+    timm_create = BACKBONES.build({'name': 'timm_resnet50'})
+    model = timm_create(model_name='resnet50')
     flops, macs, params = calculate_flops(model, (1, 3, 224, 224))
     assert flops == '8.21 GFLOPS'
     assert macs == '4.09 GMACs'