PyTorch在Ubuntu上的使用教程
在安装PyTorch前,确保Ubuntu系统软件包为最新版本,避免兼容性问题:
sudo apt update && sudo apt upgrade -y
Ubuntu通常预装Python 3,但需确认版本(建议3.6+)并安装pip(Python包管理器):
sudo apt install python3 python3-pip python3-venv -y
验证安装:
python3 --version # 查看Python版本 pip3 --version # 查看pip版本
为隔离项目依赖,避免全局包冲突,建议创建虚拟环境:
mkdir pytorch_project && cd pytorch_project python3 -m venv pytorch_env # 创建虚拟环境 source pytorch_env/bin/activate # 激活环境(激活后命令行前会显示环境名)
注:退出虚拟环境可使用
deactivate
命令。
PyTorch提供pip和conda两种主流安装方式,需根据是否需要GPU加速选择对应命令。
pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu118
注:CUDA版本需与GPU驱动兼容,可通过
nvidia-smi
命令查看驱动支持的CUDA版本。
conda install pytorch torchvision torchaudio cpuonly -c pytorch
conda install pytorch torchvision torchaudio cudatoolkit=11.7 -c pytorch -c nvidia
安装完成后,在Python交互环境中验证PyTorch是否安装成功及GPU是否可用:
import torch print("PyTorch版本:", torch.__version__) # 打印PyTorch版本 print("CUDA可用性:", torch.cuda.is_available()) # 若为True则表示GPU可用
若输出版本号且torch.cuda.is_available()
返回True
,则说明安装成功。
若安装了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()
为用户自定义的模型类,需替换为实际模型。
安装完成后,即可编写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}")
nvidia-smi
)、CUDA版本是否匹配(nvcc --version
),并重新安装对应版本的PyTorch。pip3 install torch torchvision torchaudio -i https://pypi.tuna.tsinghua.edu.cn/simple