温馨提示×

PyTorch在Ubuntu上的使用教程是什么

小樊
47
2025-09-25 14:19:50
栏目: 智能运维

PyTorch在Ubuntu上的使用教程

1. 准备工作:更新系统

在安装PyTorch前,确保Ubuntu系统软件包为最新版本,避免兼容性问题:

sudo apt update && sudo apt upgrade -y 

2. 安装Python与pip

Ubuntu通常预装Python 3,但需确认版本(建议3.6+)并安装pip(Python包管理器):

sudo apt install python3 python3-pip python3-venv -y 

验证安装:

python3 --version # 查看Python版本 pip3 --version # 查看pip版本 

3. 创建虚拟环境(可选但推荐)

为隔离项目依赖,避免全局包冲突,建议创建虚拟环境:

mkdir pytorch_project && cd pytorch_project python3 -m venv pytorch_env # 创建虚拟环境 source pytorch_env/bin/activate # 激活环境(激活后命令行前会显示环境名) 

注:退出虚拟环境可使用deactivate命令。

4. 安装PyTorch

PyTorch提供pipconda两种主流安装方式,需根据是否需要GPU加速选择对应命令。

4.1 使用pip安装

  • CPU版本(无GPU加速):
    pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu 
  • GPU版本(需NVIDIA GPU+CUDA支持):
    根据CUDA版本选择对应命令(以CUDA 11.8为例):
    pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu118 

    注:CUDA版本需与GPU驱动兼容,可通过nvidia-smi命令查看驱动支持的CUDA版本。

4.2 使用conda安装(需提前安装Anaconda/Miniconda)

  • CPU版本
    conda install pytorch torchvision torchaudio cpuonly -c pytorch 
  • GPU版本(以CUDA 11.7为例):
    conda install pytorch torchvision torchaudio cudatoolkit=11.7 -c pytorch -c nvidia 

5. 验证安装

安装完成后,在Python交互环境中验证PyTorch是否安装成功及GPU是否可用:

import torch print("PyTorch版本:", torch.__version__) # 打印PyTorch版本 print("CUDA可用性:", torch.cuda.is_available()) # 若为True则表示GPU可用 

若输出版本号且torch.cuda.is_available()返回True,则说明安装成功。

6. 配置GPU加速(可选)

若安装了GPU版本的PyTorch,需将模型与数据迁移至GPU以提升计算速度:

device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # 自动选择设备 model = Net().to(device) # 将模型迁移至设备 inputs = inputs.to(device) # 将输入数据迁移至设备 targets = targets.to(device) # 将目标数据迁移至设备 

注:Net()为用户自定义的模型类,需替换为实际模型。

7. 开始深度学习任务

安装完成后,即可编写PyTorch代码实现深度学习任务。以下是一个简单的全连接神经网络示例:

import torch import torch.nn as nn import torch.optim as optim # 定义模型 class Net(nn.Module): def __init__(self): super(Net, self).__init__() self.fc1 = nn.Linear(784, 128) # 输入层→隐藏层 self.fc2 = nn.Linear(128, 10) # 隐藏层→输出层 def forward(self, x): x = torch.flatten(x, 1) # 展平输入(如将28x28的图像展平为784维向量) x = torch.relu(self.fc1(x)) # 激活函数 x = self.fc2(x) # 输出层 return x # 初始化模型、损失函数与优化器 model = Net() criterion = nn.CrossEntropyLoss() # 交叉熵损失(适用于分类任务) optimizer = optim.SGD(model.parameters(), lr=0.01) # 随机梯度下降优化器 # 模拟输入数据(64个样本,每个样本1通道28x28像素) inputs = torch.randn(64, 1, 28, 28) targets = torch.randint(0, 10, (64,)) # 64个随机标签(0-9) # 前向传播 outputs = model(inputs) loss = criterion(outputs, targets) # 反向传播与优化 optimizer.zero_grad() # 清空梯度 loss.backward() # 反向传播计算梯度 optimizer.step() # 更新参数 print(f"训练损失: {loss.item():.4f}") 

常见问题解决

  • CUDA不可用:检查NVIDIA驱动是否安装(nvidia-smi)、CUDA版本是否匹配(nvcc --version),并重新安装对应版本的PyTorch。
  • 依赖冲突:建议使用虚拟环境隔离项目依赖,避免全局安装导致的问题。
  • 安装速度慢:可使用国内镜像源(如清华源)加速pip安装,例如:
    pip3 install torch torchvision torchaudio -i https://pypi.tuna.tsinghua.edu.cn/simple 

0