pytorch cuda安装

在安装PyTorch之前，需要确保你的机器上有NVIDIA GPU并且安装了NVIDIA CUDA工具包和cuDNN库，这是PyTorch使用GPU的前提条件。以下是在Linux系统上安装PyTorch的步骤：

1. 安装CUDA工具包和cuDNN库

首先需要安装NVIDIA CUDA Toolkit和cuDNN库。这里以Ubuntu 16.04为例：

# 安装CUDA Toolkit 10.1
$ wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1604/x86_64/cuda-ubuntu1604.pin
$ sudo mv cuda-ubuntu1604.pin /etc/apt/preferences.d/cuda-repository-pin-600
$ wget https://developer.download.nvidia.com/compute/cuda/10.1/Prod/local_installers/cuda-repo-ubuntu1604-10-1-local-10.1.105-418.39_1.0-1_amd64.deb
$ sudo dpkg -i cuda-repo-ubuntu1604-10-1-local-10.1.105-418.39_1.0-1_amd64.deb
$ sudo apt-key add /var/cuda-repo-10-1-local-10.1.105-418.39/7fa2af80.pub
$ sudo apt-get update
$ sudo apt-get install cuda

# 安装cuDNN库
$ wget https://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu1604/x86_64/libcudnn7_7.6.5.32-1+cuda10.1_amd64.deb
$ sudo dpkg -i libcudnn7_7.6.5.32-1+cuda10.1_amd64.deb
$ sudo apt-get update
$ sudo apt-get install libcudnn7

2. 安装PyTorch

安装PyTorch的最简单方法是使用conda包管理器。如果你还没有安装conda，请先安装conda。

# 创建conda环境
$ conda create --name pytorch python=3.6

# 激活conda环境
$ conda activate pytorch

# 安装PyTorch
$ conda install pytorch torchvision cudatoolkit=10.1 -c pytorch

安装完成后，可以通过以下代码测试是否安装成功：

import torch

if torch.cuda.is_available():
    device = torch.device("cuda")
    print("CUDA is available")
else:
    device = torch.device("cpu")
    print("CUDA is not available")

如果输出结果为“CUDA is available”，则说明PyTorch安装成功并且可以使用GPU加速。