[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()