如何为CUDA与ROCm硬件配置通用PyTorch虚拟环境（venv）

核心思路 #

使用未绑定特定硬件加速库的基础PyTorch包，让PyTorch在运行时自动检测目标机器的CUDA/ROCm环境并加载对应后端，实现同一个虚拟环境跨两种硬件复用。

操作步骤 #

• 创建虚拟环境

  python -m venv pytorch-unified-env
  

• 激活虚拟环境
  • Linux/macOS：

    source pytorch-unified-env/bin/activate
    

  • Windows：

    pytorch-unified-env\Scripts\activate
    

• 安装基础版PyTorch（无硬件绑定）

  pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
  

  注：此包为纯CPU基础包，但PyTorch会自动识别系统中已安装的CUDA/ROCm runtime并启用硬件加速，无需单独安装对应绑定版本。

兼容性验证 #

• 在CUDA机器上验证：

  import torch
  print(torch.cuda.is_available())  # 输出True则说明CUDA加速生效
  

• 在ROCm机器上验证：

  import torch
  print(torch.backends.rocm.is_available())  # 输出True则说明ROCm加速生效
  

关键注意事项 #

• 每台目标机器需预先安装对应硬件的系统级驱动与runtime（CUDA Toolkit或ROCm SDK），无需在虚拟环境内重复安装。
• 如需指定PyTorch版本，可修改安装命令中的版本号，例如：

  pip install torch==2.1.0 torchvision==0.16.0 torchaudio==2.1.0 --index-url https://download.pytorch.org/whl/cpu
  

• 禁止安装带有cuXXX或rocmXX.X后缀的硬件绑定版PyTorch，否则跨硬件使用会出现兼容性报错。

内容的提问来源于stack exchange，提问作者Udinanon