fix benchmark

Author: gemfield

deepvac/backbones/resnet.py

@@ -270,6 +270,7 @@ def testFly(self):
 
 def auditConfig():
     config.core.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    config.cast.ScriptCast = AttrDict()
     config.cast.ScriptCast.model_dir = "./gemfield_script.pt"
 
     config.core.disable_git = True
@@ -279,7 +280,7 @@ def auditConfig():
     config.core.log_every = 100
     config.core.num_workers = 4
 
-    config.core.net = ResNet50() # ResNet50() / resnet50()
+    config.core.net = ResNet50()
 
     config.core.optimizer = optim.SGD(
         config.core.net.parameters(),
@@ -289,6 +290,7 @@ def auditConfig():
         nesterov=False
     )
     config.core.scheduler = optim.lr_scheduler.MultiStepLR(config.core.optimizer, [50, 70, 90], 0.1)
+    config.core.criterion=torch.nn.CrossEntropyLoss()
 
     config.core.shuffle = True
     config.core.batch_size = 1
@@ -325,6 +327,7 @@ def auditConfig():
 
         config.core.val_dataset = FileLineDataset(config, fileline_path=sys.argv[5], sample_path_prefix=sys.argv[3])
         config.core.val_loader = torch.utils.data.DataLoader(config.core.val_dataset, batch_size=1, pin_memory=False)
+        config.core.test_loader = ''
         train = ResNet50Train(config)
         train()
 
@@ -333,18 +336,23 @@ def auditConfig():
             LOG.logE("Usage: python -m deepvac.backbones.resnet test <pretrained_model.pth> <your_test_img_input_dir>", exit=True)
 
         config.core.model_path = sys.argv[2]
-        config.cast.ScriptCast.model_dir = ""
+        config.cast.ScriptCast.model_dir = "./script.pt"
         config.cast.ScriptCast.static_quantize_dir = "./static_quantize.pt"
         config.core.test_dataset = ResnetClsTestDataset(config, sample_path=sys.argv[3])
         config.core.test_loader = torch.utils.data.DataLoader(config.core.test_dataset, batch_size=1, pin_memory=False)
         test = ResNet50Test(config)
         input_tensor = torch.rand(1,3,640,640)
-        test(input_tensor)
+        test()
 
     if op == 'benchmark':
         if(len(sys.argv) != 4):
             LOG.logE("Usage: python -m deepvac.backbones.resnet benchmark <pretrained_model.pth> <your_input_img.jpg>", exit=True)
 
+        config.core.model_reinterpret_cast = False
+        config.core.cast_state_dict_strict = False
+        # config.core.net_omit_keys = ['num_batches_tracked']
+        # config.core.net_omit_keys_strict = False
+        # config.core.network_audit_disabled=False
         config.core.model_path = sys.argv[2]
         config.core.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
         img_path = sys.argv[3]

deepvac/core/deepvac.py

@@ -18,7 +18,7 @@
 except ImportError:
     LOG.logE("Deepvac has dependency on tensorboard, please install tensorboard first, e.g. [pip3 install tensorboard]", exit=True)
 
-from ..utils import syszux_once, LOG, assertAndGetGitBranch, getPrintTime, AverageMeter
+from ..utils import syszux_once, LOG, assertAndGetGitBranch, getPrintTime, AverageMeter, anyFieldsInConfig
 from ..cast import export3rd
 
 #deepvac implemented based on PyTorch Framework
@@ -108,15 +108,38 @@ def initStateDict(self):
 
     def castStateDict(self, config):
         LOG.logI("model_reinterpret_cast set to True in config.py, Try to reinterpret cast the model")
+
+        if self.config.model_path_omit_keys:
+            LOG.logI("You have set config.core.model_path_omit_keys: {}".format(self.config.model_path_omit_keys))
+            for k in self.config.model_path_omit_keys:
+                LOG.logI("remove key {} from config.core.model_path {}".format(k, self.config.model_path))
+                config.state_dict.pop(k, None)
+
         state_dict = collections.OrderedDict()
         keys = list(config.state_dict.keys())
-        for idx, name in enumerate(config.net.state_dict()):
-            if config.net.state_dict()[name].size() == config.state_dict[keys[idx]].size():
-                LOG.logI("cast pretrained model [{}] => config.core.net [{}]".format(keys[idx], name))
-                state_dict[name] = config.state_dict[keys[idx]]
+        model_path_keys_len = len(keys)
+        if len(config.net.state_dict() ) > model_path_keys_len:
+            LOG.logW("config.core.net has more parameters than config.core.model_path({}), may has cast issues.".format(self.config.model_path))
+
+        real_idx = 0
+        for _, name in enumerate(config.net.state_dict()):
+            if real_idx >= model_path_keys_len:
+                LOG.logI("There alreay has no corresponding parameter in {} for {}".format(self.config.model_path, name))
+                continue
+
+            if anyFieldsInConfig(name, self.config.net_omit_keys, self.config.net_omit_keys_strict):
+                LOG.logI('found key to omit in config.core.net: {}, continue...'.format(name))
                 continue
-            LOG.logE("cannot cast pretrained model [{}] => config.net [{}] due to parameter shape mismatch!".format(keys[idx], name))
+  
+            if config.net.state_dict()[name].size() == config.state_dict[keys[real_idx]].size():
+                LOG.logI("cast pretrained model [{}] => config.core.net [{}]".format(keys[real_idx], name))
+                state_dict[name] = config.state_dict[keys[real_idx]]
+                real_idx += 1
+                continue
+
+            LOG.logE("cannot cast pretrained model [{}] => config.net [{}] due to parameter shape mismatch!".format(keys[real_idx], name))
             if config.cast_state_dict_strict is False:
+                real_idx += 1
                 continue
             LOG.logE("If you know above risk, set cast_state_dict_strict=False in config.py to omit this audit.", exit=True)
 

deepvac/utils/user_config.py

@@ -7,4 +7,17 @@ def addUserConfig(module_name, config_name, user_give=None, developer_give=None,
         return user_give
     if developer_give is not None:
         return developer_give
-    LOG.logE("value missing for configuration: {}.{} in config.py".format(module_name, config_name), exit=True)
+    LOG.logE("value missing for configuration: {}.{} in config.py".format(module_name, config_name), exit=True)
+
+def anyFieldsInConfig(name, c, c_strict=True):
+    if not c:
+        return False
+
+    if not isinstance(c, list):
+        return False
+
+    fields = name.split('.') if c_strict is False else [name]
+    for field in fields:
+        if field in c:
+            return True
+    return False