add set_random_seed and use rank shift feature (#373)

mmcv/runner/utils.py

 # Copyright (c) Open-MMLab. All rights reserved.
+import os
+import random
 import sys
 import time
 from getpass import getuser
 from socket import gethostname
 
+import numpy as np
+import torch
+
 import mmcv
 
 
@@ -48,3 +53,29 @@ def obj_from_dict(info, parent=None, default_args=None):
         for name, value in default_args.items():
             args.setdefault(name, value)
     return obj_type(**args)
+
+
+def set_random_seed(seed, deterministic=False, use_rank_shift=False):
+    """Set random seed.
+
+    Args:
+        seed (int): Seed to be used.
+        deterministic (bool): Whether to set the deterministic option for
+            CUDNN backend, i.e., set `torch.backends.cudnn.deterministic`
+            to True and `torch.backends.cudnn.benchmark` to False.
+            Default: False.
+        rank_shift (bool): Whether to add rank number to the random seed to
+            have different random seed in different threads. Default: False.
+    """
+    if use_rank_shift:
+        rank, _ = mmcv.runner.get_dist_info()
+        seed += rank
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    os.environ['PYTHONHASHSEED'] = str(seed)
+    if deterministic:
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cudnn.benchmark = False