指定GPU

方法 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