ubuntu配置pytorch gpu环境

1. 安装NVIDIA显卡驱动

首先需要安装NVIDIA显卡驱动，可以通过以下命令来安装：

sudo apt-get install nvidia-xxx

其中，xxx为显卡驱动版本号。

2. 安装CUDA

安装CUDA可以通过官网下载安装包进行安装，也可以通过以下命令进行安装：

sudo apt-get install cuda

安装完成后，需要将CUDA的路径添加到环境变量中：

export PATH=/usr/local/cuda/bin:$PATH
export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH

3. 安装cuDNN

cuDNN是NVIDIA提供的深度学习库，可以通过以下命令进行安装：

sudo apt-get install libcudnn7-dev

安装完成后，需要将cuDNN的路径添加到环境变量中：

export LD_LIBRARY_PATH=/usr/local/cuda/lib64:/usr/local/cuda/extras/CUPTI/lib64:$LD_LIBRARY_PATH

4. 安装PyTorch

可以通过pip安装PyTorch：

pip install torch torchvision

安装完成后，可以通过以下命令测试PyTorch是否能够使用GPU：

import torch
print(torch.cuda.is_available())

如果输出为True，则说明PyTorch可以使用GPU。

5. 配置PyTorch使用GPU

在使用PyTorch时，需要将模型和数据都放到GPU上进行计算。可以通过以下代码将模型和数据放到GPU上：

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
data = data.to(device)

其中，device为设备类型，可以根据需要选择使用CPU或GPU。model和data分别为模型和数据。