方法 1:在代码中设置
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
import random
import numpy as np
import torch
def seed_torch(seed=42):
random.seed(seed)
os.environ["PYTHONHASHSEED"] = str(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True
seed = 42
seed_torch(seed)
# 之后的代码
# ...
方法 2:在命令行中设置
CUDA_VISIBLE_DEVICES=0 python your_script.py --cfg your_config.yaml