55 lines
1.4 KiB
Python
55 lines
1.4 KiB
Python
# These codes are copied from modelscope revision c58451baead80d83281f063d12fb377fad415257
|
|
# Copyright (c) Alibaba, Inc. and its affiliates.
|
|
|
|
import torch
|
|
|
|
try:
|
|
import torch_xla.core.xla_model as xm
|
|
except ImportError:
|
|
xm = None
|
|
|
|
|
|
def get_rng_state():
|
|
r"""
|
|
Get random number generator state of torch, xla and cuda.
|
|
"""
|
|
state = {'torch_rng_state': torch.get_rng_state()}
|
|
if xm is not None:
|
|
state['xla_rng_state'] = xm.get_rng_state()
|
|
if torch.cuda.is_available():
|
|
state['cuda_rng_state'] = torch.cuda.get_rng_state()
|
|
return state
|
|
|
|
|
|
def set_rng_state(state):
|
|
r"""
|
|
Set random number generator state of torch, xla and cuda.
|
|
"""
|
|
torch.set_rng_state(state['torch_rng_state'])
|
|
if xm is not None:
|
|
xm.set_rng_state(state['xla_rng_state'])
|
|
if torch.cuda.is_available():
|
|
torch.cuda.set_rng_state(state['cuda_rng_state'])
|
|
|
|
|
|
class set_torch_seed(object):
|
|
r"""
|
|
Set random seed to torch, xla and cuda.
|
|
"""
|
|
|
|
def __init__(self, seed):
|
|
assert isinstance(seed, int)
|
|
self.rng_state = get_rng_state()
|
|
|
|
torch.manual_seed(seed)
|
|
if xm is not None:
|
|
xm.set_rng_state(seed)
|
|
if torch.cuda.is_available():
|
|
torch.cuda.manual_seed(seed)
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, *exc):
|
|
set_rng_state(self.rng_state)
|