跳转至

ztxexp.environment

ztxexp.environment

实验运行环境相关工具。

本模块用于： 1. 初始化深度学习实验的随机性与设备； 2. 设置当前进程优先级。

init_torch_env 

init_torch_env(seed: int = 3407, use_gpu: bool = True, gpu_id: int = 0, deterministic: bool = False, benchmark: bool = False) -> Any

初始化 PyTorch 实验环境。

参数：

名称 类型 描述 默认
seed int

全局随机种子（Python/NumPy/Torch 同步设置）。

 3407
use_gpu bool

是否优先使用 GPU。

 True
gpu_id int

当启用 GPU 时使用的设备 ID。

 0
deterministic bool

是否启用 cuDNN 确定性模式。

 False
benchmark bool

是否启用 cuDNN benchmark 自动搜索最优算法。

 False

返回：

类型 描述
Any

torch.device: 最终使用的设备对象（CPU 或 CUDA）。

引发：

类型 描述
ImportError

未安装 torch 时抛出。

示例：

>>> device = init_torch_env(seed=42, use_gpu=False)
>>> str(device)
'cpu'
源代码位于： ztxexp/environment.py 
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
def init_torch_env(
    seed: int = 3407,
    use_gpu: bool = True,
    gpu_id: int = 0,
    deterministic: bool = False,
    benchmark: bool = False,
) -> Any:
    """初始化 PyTorch 实验环境。

    Args:
        seed: 全局随机种子（Python/NumPy/Torch 同步设置）。
        use_gpu: 是否优先使用 GPU。
        gpu_id: 当启用 GPU 时使用的设备 ID。
        deterministic: 是否启用 cuDNN 确定性模式。
        benchmark: 是否启用 cuDNN benchmark 自动搜索最优算法。

    Returns:
        torch.device: 最终使用的设备对象（CPU 或 CUDA）。

    Raises:
        ImportError: 未安装 torch 时抛出。

    Examples:
        >>> device = init_torch_env(seed=42, use_gpu=False)
        >>> str(device)
        'cpu'
    """
    # 延迟导入 torch，避免非 torch 用户在导入包时失败。
    torch = _require_torch()

    # 设置三套随机种子，提升实验复现稳定性。
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)

    # GPU 路径：设置设备与 cuDNN 行为。
    if torch.cuda.is_available() and use_gpu:
        torch.cuda.manual_seed_all(seed)
        torch.cuda.set_device(gpu_id)
        torch.backends.cudnn.deterministic = deterministic
        torch.backends.cudnn.benchmark = benchmark
        print(f"Using GPU: {gpu_id}")
        return torch.device(f"cuda:{gpu_id}")

    # CPU 回退路径。
    print("Using CPU")
    return torch.device("cpu")

set_process_priority 

set_process_priority(priority: str = 'high') -> None

设置当前进程优先级。

参数：

名称 类型 描述 默认
priority str

优先级级别，可选值为 high / normal / low。

 'high'

返回：

类型 描述
None

None

引发：

类型 描述
PermissionError

当前权限不足，无法修改进程优先级。

示例：

>>> set_process_priority("normal")
源代码位于： ztxexp/environment.py 
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
def set_process_priority(priority: str = "high") -> None:
    """设置当前进程优先级。

    Args:
        priority: 优先级级别，可选值为 ``high`` / ``normal`` / ``low``。

    Returns:
        None

    Raises:
        PermissionError: 当前权限不足，无法修改进程优先级。

    Examples:
        >>> set_process_priority("normal")
    """
    process = psutil.Process(os.getpid())
    selected = priority.strip().lower()

    try:
        # Windows 使用优先级类；类 Unix 使用 nice 值。
        if os.name == "nt":
            if selected == "high":
                process.nice(psutil.HIGH_PRIORITY_CLASS)
            elif selected == "low":
                process.nice(psutil.BELOW_NORMAL_PRIORITY_CLASS)
            else:
                process.nice(psutil.NORMAL_PRIORITY_CLASS)
        else:
            if selected == "high":
                process.nice(-10)
            elif selected == "low":
                process.nice(10)
            else:
                process.nice(0)

    except (psutil.AccessDenied, PermissionError) as exc:  # pragma: no cover
        raise PermissionError(
            "Insufficient permission to change process priority. "
            "Try running with elevated privileges or use 'normal'."
        ) from exc