update to suit new timm models

.gitignore

@@ -2,6 +2,7 @@
 /ckpts
 /data/voc
 /wandb
+/report
 
 # Pycharm
 .idea

requirements.txt

@@ -1,7 +1,7 @@
 torch~=1.7.1
 torchvision~=0.8.2
 einops~=0.3.0
-timm~=0.3.2
+timm~=0.4.12
 numpy~=1.18.5
 tqdm~=4.54.0
 wandb~=0.10.23

utils.py

@@ -6,6 +6,7 @@
 import torch
 import torchvision
 import torchvision.models as models
+import timm
 from timm.models import create_model
 from torch.utils.data import DataLoader
 from torchvision import transforms
@@ -45,6 +46,7 @@ def get_model(args, pretrained=True):
     model_names = sorted(name for name in models.__dict__
                          if name.islower() and not name.startswith("__")
                          and callable(models.__dict__[name]))
+    timm_model_names = timm.list_models(pretrained=True)
 
     if args.model_name in model_names:
         model = models.__dict__[args.model_name](pretrained=pretrained)
@@ -86,7 +88,7 @@ def get_model(args, pretrained=True):
         model = create_model(args.model_name, pretrained=pretrained)
         mean = (0.5, 0.5, 0.5)
         std = (0.5, 0.5, 0.5)
-    elif 'regnety' in args.model_name:
+    elif args.model_name in timm_model_names:
         model = create_model(args.model_name, pretrained=pretrained)
         mean = (0.485, 0.456, 0.406)
         std = (0.229, 0.224, 0.225)