mirror of https://github.com/open-mmlab/mmcv.git
39 lines
1.2 KiB
Python
39 lines
1.2 KiB
Python
import os
|
|
import random
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
from mmcv.runner import set_random_seed
|
|
from mmcv.utils import TORCH_VERSION, digit_version
|
|
|
|
is_rocm_pytorch = False
|
|
if digit_version(TORCH_VERSION) >= digit_version('1.5'):
|
|
from torch.utils.cpp_extension import ROCM_HOME
|
|
is_rocm_pytorch = True if ((torch.version.hip is not None) and
|
|
(ROCM_HOME is not None)) else False
|
|
|
|
|
|
def test_set_random_seed():
|
|
set_random_seed(0)
|
|
a_random = random.randint(0, 10)
|
|
a_np_random = np.random.rand(2, 2)
|
|
a_torch_random = torch.rand(2, 2)
|
|
assert torch.backends.cudnn.deterministic is False
|
|
assert torch.backends.cudnn.benchmark is False
|
|
assert os.environ['PYTHONHASHSEED'] == str(0)
|
|
|
|
set_random_seed(0, True)
|
|
b_random = random.randint(0, 10)
|
|
b_np_random = np.random.rand(2, 2)
|
|
b_torch_random = torch.rand(2, 2)
|
|
assert torch.backends.cudnn.deterministic is True
|
|
if is_rocm_pytorch:
|
|
assert torch.backends.cudnn.benchmark is True
|
|
else:
|
|
assert torch.backends.cudnn.benchmark is False
|
|
|
|
assert a_random == b_random
|
|
assert np.equal(a_np_random, b_np_random).all()
|
|
assert torch.equal(a_torch_random, b_torch_random)
|