根治 PyTorch CUDA `pynvml` 弃用警告:直接修改 `torch/cuda/__init__.py` 的实践记录

发布时间:2026/6/5 1:16:03

根治 PyTorch CUDA `pynvml` 弃用警告:直接修改 `torch/cuda/__init__.py` 的实践记录 根治 PyTorch CUDApynvml弃用警告直接修改torch/cuda/__init__.py的实践记录【笔记】消除 “FutureWarning: The pynvml package is deprecated”警告1. 问题现象在 Windows 环境下运行 ComfyUI或任何基于 PyTorch 的 CUDA 应用时控制台会在 PyTorch 初始化阶段抛出一条FutureWarningH:\PythonProjects3\Win_ComfyUI\.venv\Lib\site-packages\torch\cuda\__init__.py:61: FutureWarning: The pynvml package is deprecated. Please install nvidia-ml-py instead. import pynvml # type: ignore[import]虽然这条警告不影响 CUDA 的正常使用也不会阻断程序运行但在每次启动时污染日志输出尤其在使用--verbose或排查其他问题时显得格外扎眼。2. 根因分析这条警告的源头并非 PyTorch 本身而是pynvml这个包的设计策略。pynvml13.x 的本质当前 PyPI 上的pynvml已经演变成一个兼容 shim垫片。它内部完全依赖nvidia-ml-py提供 NVMLNVIDIA Management Library的 Python 绑定但会在被import时主动抛出FutureWarning提醒用户迁移。PyTorch 的行为PyTorch 在torch.cuda初始化时硬编码了import pynvml来获取驱动版本、显存信息等。由于它先尝试导入pynvml自然就触发了 shim 层的警告。环境的现状通过pipdeptree可以确认nvidia-ml-py已经正确安装且被flashinfer-python、inference等包直接依赖而pynvml的Required-by为空说明没有硬依赖它的上层包。因此真正提供功能的是nvidia-ml-py而pynvml只是一个多此一举的警告发射器。pipdeptree--reverse--packages pynvml3. 解决思路让 PyTorch 直接对接底层库与其用环境变量或warnings.filterwarnings去“捂嘴”屏蔽警告不如从源头根治让 PyTorch 跳过pynvml这个 shim直接导入nvidia-ml-py。由于 PyTorch 后续代码中大量使用了pynvml.nvmlInit()、pynvml.nvmlDeviceGetCount()等调用我们只需要在导入时做一个别名映射让pynvml这个命名空间直接指向nvidia-ml-py即可。这样无需改动 PyTorch 的任何业务逻辑实现零侵入式修复。4. 具体操作步骤定位文件请根据你的实际虚拟环境路径调整H:\PythonProjects3\Win_ComfyUI\.venv\Lib\site-packages\torch\cuda\__init__.py找到对应代码块大约在第 59~61 行try:ifnot_version.hip:importpynvml# type: ignore[import]修改后的代码try:ifnot_version.hip:# import pynvml # type: ignore[import]importnvidia_ml_pyaspynvml关键细节保持缩进确保新添加的import nvidia_ml_py as pynvml与原import pynvml的缩进完全一致4 个空格否则 Python 解释器会抛出IndentationError。注释而非删除将原行注释掉而不是直接删除。这样一旦需要回滚只需取消注释即可无需记忆修改内容。添加书签在编辑器如 VS Code、PyCharm、Notepad 等中为该行添加书签。PyTorch 重装pip install torch2.7.1 torchvision0.22.1 torchaudio2.7.1 --index-url https://download.pytorch.org/whl/cu126 --force-reinstall会覆盖此文件书签能帮你下次快速定位并重新修改。pip install torch2.7.1torchvision0.22.1torchaudio2.7.1--index-url https://download.pytorch.org/whl/cu126--force-reinstall5. 验证效果保存文件后重启 ComfyUI。观察日志修改前FutureWarning: The pynvml package is deprecated...出现在torch/cuda/__init__.py:61修改后该位置警告完全消失且 CUDA 初始化、GPU 检测、显存查询等功能一切正常6. 注意事项与回溯策略事项说明虚拟环境隔离修改仅作用于当前.venv不会影响系统 Python 或其他项目的 PyTorch 安装PyTorch 重装执行torch重装 后__init__.py会被覆盖需重新执行上述修改。这也是建议添加书签的原因版本兼容性nvidia-ml-py是 NVML 的官方 Python 绑定API 与旧pynvml完全兼容不存在功能风险备份建议如果你担心改坏可先复制一份__init__.py为__init__.py.bak7. 与其他方案的对比为什么推荐直接改源码而不是其他更“温和”的手段方案缺点set PYTHONWARNINGSignore::FutureWarning过于粗暴会同时屏蔽项目中其他可能有价值的FutureWarning在sitecustomize.py中warnings.filterwarnings需要额外维护一个文件且属于“掩耳盗铃”警告的触发逻辑仍在运行直接pip uninstall pynvml虽然pipdeptree显示无依赖但某些 ComfyUI 自定义节点可能硬编码import pynvml卸载后可能导致其报错相比之下修改torch/cuda/__init__.py的别名导入是风险最低、最优雅的根治方案它保留了pynvml包在环境中的存在安抚可能依赖它的第三方插件同时让 PyTorch 直接走官方库彻底消除警告。8. 总结这是一次典型的“面向源码排障”实践。通过分析依赖树的实际指向我们发现pynvml已经退化为无实质功能的警告层而 PyTorch 恰好是这个警告的触发者。利用 Pythonimport ... as ...的别名机制我们在 PyTorch 内部做了一个无缝的“拨乱反正”既保留了环境兼容性又实现了日志的彻底清净。下次重装 / 升级 PyTorch 后只需凭借编辑器书签回到同一位置重复一次注释与替换即可继续享受零警告的启动体验。

相关新闻