githubmirror/ColossalAI

colossalai/testing/random.py

19 строк544 B

HELSON

[zero] test gradient accumulation (#1964)

29 ноя 2022, 08:00

29 ноя 2022, 08:00a1ce02d

import random

import numpy as np
import torch


def seed_all(seed, cuda_deterministic=False):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
    if cuda_deterministic:    # slower, more reproducible
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False
    else:
        torch.backends.cudnn.deterministic = False
        torch.backends.cudnn.benchmark = True