[Fix] Fix iter bug when resuming checkpoint in distributed train (open-mmlab#866)

* [Fix] Fix iter bug when resuming checkpoint in distributed train

* fix lint error

Signed-off-by: FreyWang <[email protected]>

Author: FreyWang

tools/train.py

@@ -7,7 +7,7 @@
 
 import mmcv
 import torch
-from mmcv.runner import init_dist
+from mmcv.runner import get_dist_info, init_dist
 from mmcv.utils import Config, DictAction, get_git_hash
 
 from mmseg import __version__
@@ -94,6 +94,9 @@ def main():
     else:
         distributed = True
         init_dist(args.launcher, **cfg.dist_params)
+        # gpu_ids is used to calculate iter when resuming checkpoint,
+        _, world_size = get_dist_info()
+        cfg.gpu_ids = range(world_size)
 
     # create work_dir
     mmcv.mkdir_or_exist(osp.abspath(cfg.work_dir))