diff --git a/mmrazor/apis/__init__.py b/mmrazor/apis/__init__.py index 3da568f6..bc7b8f48 100644 --- a/mmrazor/apis/__init__.py +++ b/mmrazor/apis/__init__.py @@ -2,4 +2,6 @@ from .mmcls import * # noqa: F401,F403 from .mmdet import * # noqa: F401,F403 from .mmseg import * # noqa: F401,F403 -from .utils import set_random_seed # noqa: F401 +from .utils import init_random_seed, set_random_seed # noqa: F401 + +__all__ = ['init_random_seed', 'set_random_seed'] diff --git a/mmrazor/apis/utils.py b/mmrazor/apis/utils.py index bdc6d767..e00c5780 100644 --- a/mmrazor/apis/utils.py +++ b/mmrazor/apis/utils.py @@ -3,6 +3,39 @@ import random import numpy as np import torch +import torch.distributed as dist +from mmcv.runner import get_dist_info + + +def init_random_seed(seed=None, device='cuda'): + """Initialize random seed. + + If the seed is not set, the seed will be automatically randomized, + and then broadcast to all processes to prevent some potential bugs. + Args: + seed (int, Optional): The seed. Default to None. + device (str): The device where the seed will be put on. + Default to 'cuda'. + Returns: + int: Seed to be used. + """ + if seed is not None: + return seed + + # Make sure all ranks share the same random seed to prevent + # some potential bugs. Please refer to + # https://github.com/open-mmlab/mmdetection/issues/6339 + rank, world_size = get_dist_info() + seed = np.random.randint(2**31) + if world_size == 1: + return seed + + if rank == 0: + random_num = torch.tensor(seed, dtype=torch.int32, device=device) + else: + random_num = torch.tensor(0, dtype=torch.int32, device=device) + dist.broadcast(random_num, src=0) + return random_num.item() def set_random_seed(seed: int, deterministic: bool = False) -> None: diff --git a/tools/mmcls/train_mmcls.py b/tools/mmcls/train_mmcls.py index a5d6144d..43034750 100644 --- a/tools/mmcls/train_mmcls.py +++ b/tools/mmcls/train_mmcls.py @@ -8,6 +8,7 @@ import warnings import mmcv import torch +import torch.distributed as dist from mmcls import __version__ from mmcls.datasets import build_dataset from mmcls.utils import collect_env, get_root_logger @@ -15,7 +16,7 @@ from mmcv import Config, DictAction from mmcv.runner import get_dist_info, init_dist # Differences from mmclassification -from mmrazor.apis import set_random_seed, train_mmcls_model +from mmrazor.apis import init_random_seed, set_random_seed, train_mmcls_model from mmrazor.models import build_algorithm from mmrazor.utils import setup_multi_processes @@ -54,6 +55,10 @@ def parse_args(): help='id of gpu to use ' '(only applicable to non-distributed training)') parser.add_argument('--seed', type=int, default=None, help='random seed') + parser.add_argument( + '--diff_seed', + action='store_true', + help='Whether or not set different seeds for different ranks') parser.add_argument( '--deterministic', action='store_true', @@ -154,12 +159,14 @@ def main(): logger.info(f'Config:\n{cfg.pretty_text}') # set random seeds - if args.seed is not None: - logger.info(f'Set random seed to {args.seed}, ' - f'deterministic: {args.deterministic}') - set_random_seed(args.seed, deterministic=args.deterministic) - cfg.seed = args.seed - meta['seed'] = args.seed + seed = init_random_seed(args.seed) + seed = seed + dist.get_rank() if args.diff_seed else seed + logger.info(f'Set random seed to {seed}, ' + f'deterministic: {args.deterministic}') + set_random_seed(seed, deterministic=args.deterministic) + cfg.seed = seed + meta['seed'] = seed + meta['exp_name'] = osp.basename(args.config) # Difference from mmclassification # replace `model` to `algorithm` diff --git a/tools/mmdet/train_mmdet.py b/tools/mmdet/train_mmdet.py index 0c2ef363..688f8673 100644 --- a/tools/mmdet/train_mmdet.py +++ b/tools/mmdet/train_mmdet.py @@ -16,6 +16,7 @@ import warnings import mmcv import torch +import torch.distributed as dist from mmcv import Config, DictAction from mmcv.runner import get_dist_info, init_dist from mmcv.utils import get_git_hash @@ -23,7 +24,7 @@ from mmdet import __version__ from mmdet.datasets import build_dataset from mmdet.utils import collect_env, get_root_logger -from mmrazor.apis import set_random_seed, train_mmdet_model +from mmrazor.apis import init_random_seed, set_random_seed, train_mmdet_model from mmrazor.models import build_algorithm from mmrazor.utils import setup_multi_processes @@ -61,6 +62,10 @@ def parse_args(): help='id of gpu to use ' '(only applicable to non-distributed training)') parser.add_argument('--seed', type=int, default=None, help='random seed') + parser.add_argument( + '--diff_seed', + action='store_true', + help='Whether or not set different seeds for different ranks') parser.add_argument( '--deterministic', action='store_true', @@ -166,12 +171,13 @@ def main(): logger.info(f'Config:\n{cfg.pretty_text}') # set random seeds - if args.seed is not None: - logger.info(f'Set random seed to {args.seed}, ' - f'deterministic: {args.deterministic}') - set_random_seed(args.seed, deterministic=args.deterministic) - cfg.seed = args.seed - meta['seed'] = args.seed + seed = init_random_seed(args.seed) + seed = seed + dist.get_rank() if args.diff_seed else seed + logger.info(f'Set random seed to {seed}, ' + f'deterministic: {args.deterministic}') + set_random_seed(seed, deterministic=args.deterministic) + cfg.seed = seed + meta['seed'] = seed meta['exp_name'] = osp.basename(args.config) algorithm = build_algorithm(cfg.algorithm) diff --git a/tools/mmseg/train_mmseg.py b/tools/mmseg/train_mmseg.py index a32eab45..d80978d9 100644 --- a/tools/mmseg/train_mmseg.py +++ b/tools/mmseg/train_mmseg.py @@ -16,6 +16,7 @@ import warnings import mmcv import torch +import torch.distributed as dist from mmcv.cnn.utils import revert_sync_batchnorm from mmcv.runner import get_dist_info, init_dist from mmcv.utils import Config, DictAction, get_git_hash @@ -24,7 +25,7 @@ from mmseg.datasets import build_dataset from mmseg.utils import collect_env, get_root_logger # Differences from mmdetection -from mmrazor.apis import set_random_seed, train_mmseg_model +from mmrazor.apis import init_random_seed, set_random_seed, train_mmseg_model from mmrazor.models.builder import build_algorithm from mmrazor.utils import setup_multi_processes @@ -64,6 +65,10 @@ def parse_args(): help='id of gpu to use ' '(only applicable to non-distributed training)') parser.add_argument('--seed', type=int, default=None, help='random seed') + parser.add_argument( + '--diff_seed', + action='store_true', + help='Whether or not set different seeds for different ranks') parser.add_argument( '--deterministic', action='store_true', @@ -168,11 +173,12 @@ def main(): logger.info(f'Config:\n{cfg.pretty_text}') # set random seeds - if args.seed is not None: - logger.info(f'Set random seed to {args.seed}, deterministic: ' - f'{args.deterministic}') - set_random_seed(args.seed, deterministic=args.deterministic) - cfg.seed = args.seed + seed = init_random_seed(args.seed) + seed = seed + dist.get_rank() if args.diff_seed else seed + logger.info(f'Set random seed to {args.seed}, deterministic: ' + f'{args.deterministic}') + set_random_seed(args.seed, deterministic=args.deterministic) + cfg.seed = seed meta['seed'] = args.seed meta['exp_name'] = osp.basename(args.config)