[Enhancement] Support flexible runner (#2646)

open-mmlab · Sep 6, 2023 · de82f14 · de82f14
1 parent 8ff889a
commit de82f14
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 2 deletions.
diff --git a/tools/test.py b/tools/test.py
@@ -6,6 +6,8 @@
 from mmengine.config import Config, DictAction
 from mmengine.runner import Runner
 
+from mmaction.registry import RUNNERS
+
 
 def parse_args():
     parser = argparse.ArgumentParser(
@@ -108,7 +110,13 @@ def main():
     cfg.load_from = args.checkpoint
 
     # build the runner from config
-    runner = Runner.from_cfg(cfg)
+    if 'runner_type' not in cfg:
+        # build the default runner
+        runner = Runner.from_cfg(cfg)
+    else:
+        # build customized runner from the registry
+        # if 'runner_type' is set in the cfg
+        runner = RUNNERS.build(cfg)
 
     # start testing
     runner.test()

diff --git a/tools/train.py b/tools/train.py
@@ -6,6 +6,8 @@
 from mmengine.config import Config, DictAction
 from mmengine.runner import Runner
 
+from mmaction.registry import RUNNERS
+
 
 def parse_args():
     parser = argparse.ArgumentParser(description='Train a action recognizer')
@@ -125,7 +127,13 @@ def main():
     cfg = merge_args(cfg, args)
 
     # build the runner from config
-    runner = Runner.from_cfg(cfg)
+    if 'runner_type' not in cfg:
+        # build the default runner
+        runner = Runner.from_cfg(cfg)
+    else:
+        # build customized runner from the registry
+        # if 'runner_type' is set in the cfg
+        runner = RUNNERS.build(cfg)
 
     # start training
     runner.train()