File size: 1,490 Bytes
9e14838 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 | #!/usr/bin/env python
"""
startup_config
Startup configuration utilities
"""
from __future__ import absolute_import
import importlib
import os
import random
import sys
import numpy as np
import torch
__author__ = "Xin Wang"
__email__ = "wangxin@nii.ac.jp"
__copyright__ = "Copyright 2020, Xin Wang"
def set_random_seed(random_seed, args=None):
"""set_random_seed(random_seed, args=None)
Set the random_seed for numpy, python, and cudnn
input
-----
random_seed: integer random seed
args: argue parser
"""
# initialization
torch.manual_seed(random_seed)
random.seed(random_seed)
np.random.seed(random_seed)
os.environ["PYTHONHASHSEED"] = str(random_seed)
# For torch.backends.cudnn.deterministic
# Note: this default configuration may result in RuntimeError
# see https://pytorch.org/docs/stable/notes/randomness.html
if args is None:
cudnn_deterministic = True
cudnn_benchmark = False
else:
cudnn_deterministic = args.cudnn_deterministic_toggle
cudnn_benchmark = args.cudnn_benchmark_toggle
if not cudnn_deterministic:
print("cudnn_deterministic set to False")
if cudnn_benchmark:
print("cudnn_benchmark set to True")
if torch.cuda.is_available():
torch.cuda.manual_seed_all(random_seed)
torch.backends.cudnn.deterministic = cudnn_deterministic
torch.backends.cudnn.benchmark = cudnn_benchmark
return
|