29 lines
610 B
Python
29 lines
610 B
Python
|
import random
|
|||
|
|
|
|||
|
|
import numpy as np
|
|||
|
|
|
|||
|
|
from .data import *
|
|||
|
|
from .metrics import *
|
|||
|
|
|
|||
|
|
|
|||
|
|
def set_random_seed(seed):
|
|||
|
|
"""设置Python, numpy, PyTorch的随机数种子
|
|||
|
|
|
|||
|
|
:param seed: int 随机数种子
|
|||
|
|
"""
|
|||
|
|
random.seed(seed)
|
|||
|
|
np.random.seed(seed)
|
|||
|
|
torch.manual_seed(seed)
|
|||
|
|
if torch.cuda.is_available():
|
|||
|
|
torch.cuda.manual_seed(seed)
|
|||
|
|
dgl.seed(seed)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def get_device(device):
|
|||
|
|
"""返回指定的GPU设备
|
|||
|
|
|
|||
|
|
:param device: int GPU编号，-1表示CPU
|
|||
|
|
:return: torch.device
|
|||
|
|
"""
|
|||
|
|
return torch.device(f'cuda:{device}' if device >= 0 and torch.cuda.is_available() else 'cpu')
|