Implement for convnext trainer

VIAME · Sep 8, 2023 · 090a103 · 090a103
1 parent bbd5694
commit 090a103
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/plugins/pytorch/cutler/cutler_trainer.py b/plugins/pytorch/cutler/cutler_trainer.py
@@ -78,6 +78,8 @@ class ConvNextCascadeRCNNTrainer( TrainDetector ):
 
         _Option('_cutler_config_file', 'cutler_config_file', '', str, ''),
 
+        _Option('_seed_weights', 'seed_weights', '', str, ''),
+
         _Option('_output_directory', 'output_directory', 'category_models', str, ''),
         _Option('_pipeline_template', 'pipeline_template', '', str, '')
     ]
@@ -245,7 +247,9 @@ def set_configuration( self, cfg_in ):
                 os.environ[ "MASTER_PORT" ] = "12345"
             init_dist( self._launcher )
 
-        if "checkpoint_override" in self.config and self.config["checkpoint_override"]:
+        if self._seed_weights:
+            self.original_chkpt_file = self._seed_weights
+        elif "checkpoint_override" in self.config and self.config["checkpoint_override"]:
             self.original_chkpt_file = self.config["checkpoint_override"]
         else:
             self.original_chkpt_file = self.config["model_checkpoint_file"]

diff --git a/tools/viame_train_detector.cxx b/tools/viame_train_detector.cxx
@@ -1297,8 +1297,8 @@ main( int argc, char* argv[] )
     std::map< std::string, std::vector< std::string > > weight_ext =
       {
         { ".zip", { "seed_model" } },
-        { ".pth", { "backbone" } },
-        { ".pt", { "backbone" } },
+        { ".pth", { "backbone", "seed_weights" } },
+        { ".pt", { "backbone", "seed_weights" } },
         { ".py", { "config" } },
         { ".weights", { "seed_weights" } },
         { ".wt", { "seed_weights" } }