Ubuntu 20下不同PyTorch版本虚拟环境运行速度差异排查求助

可能的核心原因 #

1. PyTorch与CUDA/CuDNN版本不匹配
   PyTorch 1.7.1+cu110发布时，官方适配的CuDNN版本为8.0.x，而你当前使用的CuDNN 8.7.0是为更高版本PyTorch/CUDA设计的，高版本CuDNN与旧PyTorch的兼容性存在问题，可能导致部分CUDA算子 fallback 到低效实现甚至CPU计算。

2. Ampere架构支持不足
   RTX 3060属于Ampere架构（算力8.6），PyTorch 1.7.1对Ampere的支持仅为实验性，默认未完全开启Tensor Core等硬件优化；而PyTorch 1.11.0已经对Ampere架构做了全面优化，算子计算效率大幅提升。

3. 旧环境缺失优化依赖库
   新环境中额外安装的库可能包含优化加速组件，比如带MKL加速的numpy、pillow-simd（图像预处理加速）、torchvision的CUDA优化版本等，旧环境缺少这些依赖会导致数据预处理或辅助计算拖慢整体迭代速度。

4. 环境变量配置异常
   旧环境的LD_LIBRARY_PATH可能未正确指向CUDA 11.0的库文件，或混入了其他版本的CUDA路径，导致PyTorch加载CUDA库时出现异常，无法使用硬件加速。

解决步骤 #

1. 验证CUDA/CuDNN在旧环境的加载状态 #

在旧虚拟环境中执行以下命令，确认基础组件是否正常工作：

python -c "
import torch
print('CUDA可用:', torch.cuda.is_available())
print('当前CuDNN版本:', torch.backends.cudnn.version())
print('GPU设备名:', torch.cuda.get_device_name(0))
print('GPU算力:', torch.cuda.get_device_capability(0))
"

• 如果CuDNN版本显示不是8.0.x，说明加载了不匹配的高版本CuDNN，需要替换为PyTorch 1.7.1适配的版本。
• 如果CUDA不可用或设备名异常，检查环境变量配置。

2. 替换匹配的CuDNN版本 #

下载PyTorch 1.7.1+cu110对应的CuDNN 8.0.5（CUDA 11.0版本），解压后替换系统中对应的库文件：

# 备份当前CuDNN库
sudo cp /usr/local/cuda/include/cudnn*.h /usr/local/cuda/include/cudnn_backup/
sudo cp /usr/local/cuda/lib64/libcudnn*.so* /usr/local/cuda/lib64/cudnn_backup/

# 解压下载的CuDNN包，复制到CUDA目录
sudo cp include/cudnn*.h /usr/local/cuda/include/
sudo cp lib64/libcudnn*.so* /usr/local/cuda/lib64/
sudo chmod a+r /usr/local/cuda/include/cudnn*.h /usr/local/cuda/lib64/libcudnn*.so*

替换后重启虚拟环境，重新验证CuDNN版本。

3. 补全旧环境的优化依赖 #

对比新环境的依赖列表（执行pip freeze > new_env_requirements.txt），将旧环境缺失的优化相关库安装到位：

• 确保安装对应版本的torchvision（PyTorch 1.7.1对应torchvision==0.8.2+cu110）：

  pip install torchvision==0.8.2+cu110 -f https://download.pytorch.org/whl/torch_stable.html
  

• 安装带MKL加速的numpy：

  pip install numpy --upgrade --force-reinstall
  

• 如果AutoNovel涉及图像预处理，安装pillow-simd替代默认pillow：

  pip uninstall pillow -y
  pip install pillow-simd
  

4. 修复环境变量配置 #

激活旧虚拟环境后，检查并设置正确的CUDA环境变量：

# 清除可能的冲突路径
export LD_LIBRARY_PATH=/usr/local/cuda-11.0/lib64:$LD_LIBRARY_PATH
# 指定仅使用第一块GPU
export CUDA_VISIBLE_DEVICES=0

可以将上述命令添加到虚拟环境的activate脚本中，确保每次激活自动配置。

5. 开启Ampere架构优化（实验性） #

由于PyTorch 1.7.1对Ampere的支持有限，可尝试设置环境变量强制开启架构优化：

export TORCH_CUDA_ARCH_LIST="8.6"

然后重新运行AutoNovel代码，观察速度是否提升。

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