move metrics - support url for loading ckpt

Author: YassineYousfi

LitModel.py

@@ -8,7 +8,8 @@
 from optim.schedulers import get_lr_scheduler
 import models.surgeries
 from models.models import get_net
-from metrics import Accuracy, wAUC, PE, MD5
+from metrics.roc_metrics import wAUC, PE, MD5
+from torchmetrics.classification.accuracy import Accuracy
 
 class LitModel(pl.LightningModule):
     """
@@ -41,7 +42,7 @@ def __build_model(self):
         self.net = get_net(self.args.model.backbone, 
                            num_classes=self.num_classes,
                            in_chans=self.in_chans,
-                           imagenet=self.args.ckpt.imagenet, 
+                           pretrained=self.args.ckpt.pretrained, 
                            ckpt_path=self.args.ckpt.seed_from)
 
         # 2. Do surgery if needed

cfg/default.yaml

@@ -19,7 +19,7 @@ ckpt:
   resume_from: null
   seed_from: null
   load_fc: True
-  imagenet: True
+  pretrained: True
 
 optimizer:
   # do not decay batch norm and bias and FC

metrics.py metrics/roc_metrics.pymetrics.py renamed to metrics/roc_metrics.py

@@ -11,7 +11,6 @@
 from tools.numpy_utils import check_nans
 from torchmetrics.utilities.data import dim_zero_cat
 from torchmetrics.classification.auc import AUC
-from torchmetrics.classification.accuracy import Accuracy
 from torchmetrics.functional.classification.auc import _auc_update
 
 @_numpy_metric_conversion

models/models.py

@@ -1,6 +1,7 @@
 import warnings
 warnings.simplefilter(action='ignore', category=FutureWarning)
 from functools import partial
+from validators import url
 import numpy as np
 from torch import nn
 import timm
@@ -38,12 +39,15 @@
 
 }
 
-def get_net(model_name, num_classes=2, in_chans=3, imagenet=True, ckpt_path=None, strict_loading=False):
-    net = zoo_params[model_name]['init_op'](num_classes=num_classes, in_chans=in_chans, pretrained=imagenet)
+def get_net(model_name, num_classes=2, in_chans=3, pretrained=True, ckpt_path=None, strict_loading=False):
+    net = zoo_params[model_name]['init_op'](num_classes=num_classes, in_chans=in_chans, pretrained=pretrained)
     net.model_name = model_name
-
+    
     if ckpt_path is not None:
-        state_dict = torch.load(ckpt_path)['state_dict']
+        if url(ckpt_path):
+            state_dict = torch.hub.load_state_dict_from_url(ckpt_path)['state_dict']
+        else:
+            state_dict = torch.load(ckpt_path)['state_dict']
         state_dict = {k.split('net.')[1]: v for k, v in state_dict.items()}
 
         # Check FC compatibility

requirements.txt

@@ -13,3 +13,4 @@ omegaconf==2.1.1
 wonderwords==2.2.0
 braceexpand==0.1.7
 wandb
+validators

tests/test_config.yaml

@@ -19,7 +19,7 @@ ckpt:
   resume_from: null
   seed_from: null
   load_fc: True
-  imagenet: True
+  pretrained: True
 
 optimizer:
   # do not decay batch norm and bias and FC