一文搞懂PyTorch与CUDA那些事

PyTorch的CUDA版本和电脑/服务器自身的CUDA版本没有关系！！！

甚至和电脑/服务器有没有CUDA都没有关系！！！

许多同学误以为PyTorch和TensorFlow一样，运行时需要本机环境内的CUDA和CuDNN支持，但其实在一台没有安装CUDA的环境中按照官网的指南安装PyTorch也是可以正常在GPU上运行的。那么安装GPU版本的PyTorch需要什么？答案是：GPU和驱动。

我们看一下官网给出的安装教程：

image.png

在选择不同CUDA版本时，PyTorch版本号尾部会出现cuxxx的字样，也就是说，PyTorch在安装时会自动安装所需的CUDA运行库，用户只需要保持PyTorch版本（包含其内部的CUDA版本）、Nvidia驱动版本和GPU型号相匹配。

CUDA版本与驱动版本对应关系见下图：

image.png

通常情况下，只要驱动满足要求，并且PyTorch官方有预编译相应的CUDA版本，直接安装对应版本的PyTorch即可。

注意：安培（Ampere）架构GPU只支持CUDA11.0即以上版本，PyTorch官方只有1.7及以上版本才有预编译CUDA11，因此建议使用安培架构GPU的同学使用PyTorch1.7及以上版本