Fix mmdet

mmtune/mm/tasks/mmdet.py

@@ -1,5 +1,6 @@
 import argparse
 import copy
+import os
 import time
 from os import path as osp
 from typing import Optional, Sequence
@@ -59,11 +60,14 @@ def parse_args(self, args: Sequence[str]) -> argparse.Namespace:
             'like key="[a,b]" or key=a,b It also allows nested list/tuple '
             'values, e.g. key="[(a,b),(c,d)]" Note that the quotation marks '
             'are necessary and that no white space is allowed.')
+        parser.add_argument('--local_rank', type=int, default=0)
         parser.add_argument(
             '--auto-scale-lr',
             action='store_true',
             help='enable automatically scaling LR.')
         args = parser.parse_args(args)
+        if 'LOCAL_RANK' not in os.environ:
+            os.environ['LOCAL_RANK'] = str(args.local_rank)
         return args
 
     def build_model(self,
@@ -95,13 +99,14 @@ def train_model(self,
     def run(self, *args, **kwargs):
         from mmdet import __version__
         from mmdet.apis import init_random_seed, set_random_seed
-        from mmdet.utils import (collect_env, get_root_logger,
+        from mmdet.utils import (collect_env, get_device, get_root_logger,
                                  setup_multi_processes)
         args = self.args
 
         cfg = Config.fromfile(args.config)
         if args.cfg_options is not None:
             cfg.merge_from_dict(args.cfg_options)
+        cfg.device = ''
 
         # set cudnn_benchmark
         if cfg.get('cudnn_benchmark', False):
@@ -155,8 +160,9 @@ def run(self, *args, **kwargs):
         logger.info(f'Distributed training: {distributed}')
         logger.info(f'Config:\n{cfg.pretty_text}')
 
+        cfg.device = get_device()
         # set random seeds
-        seed = init_random_seed(args.seed)
+        seed = init_random_seed(args.seed, device=cfg.device)
         seed = seed + dist.get_rank() if args.diff_seed else seed
         logger.info(f'Set random seed to {seed}, '
                     f'deterministic: {args.deterministic}')