在Ubuntu上使用PyTorch前，需先安装系统级依赖和开发工具，确保后续步骤顺利进行。
sudo apt update && sudo apt upgrade -y，同步系统软件包列表并升级现有包。sudo apt install python3 python3-pip python3-venv -y安装Python解释器和包管理工具（venv用于创建虚拟环境）。sudo apt install -y build-essential cmake git wget unzip yasm pkg-config libopenblas-dev liblapack-dev libjpeg-dev libpng-dev，这些工具用于编译PyTorch扩展或处理图像数据。PyTorch支持CPU和GPU加速（需NVIDIA GPU及CUDA支持），推荐通过pip或conda安装预编译版本。
python3 -m venv pytorch_env # 创建名为pytorch_env的虚拟环境
source pytorch_env/bin/activate # 激活环境（终端提示符会显示环境名）
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpupip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu118
注：CUDA版本需与GPU驱动兼容（可通过
nvidia-smi查看驱动支持的CUDA版本）。
打开Python终端，运行以下代码验证安装：
import torch
print("PyTorch版本:", torch.__version__) # 输出版本号（如2.1.0）
print("CUDA可用性:", torch.cuda.is_available()) # GPU版本应返回True
若输出版本号且torch.cuda.is_available()为True，则安装成功。
MNIST是手写数字图像数据集（28x28像素），适合入门图像分类任务。
import torch
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
# 定义数据预处理：转换为张量并归一化（均值0.1307，标准差0.3081）
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])
# 加载训练集（下载到./data目录）
train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True) # 批大小64，打乱数据
# 加载测试集
test_dataset = datasets.MNIST(root='./data', train=False, transform=transform)
test_loader = DataLoader(test_dataset, batch_size=1000, shuffle=False)
CNN适合处理图像数据，通过卷积层提取特征。
import torch.nn as nn
import torch.nn.functional as F
class CNN(nn.Module):
def __init__(self):
super(CNN, self).__init__()
self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1) # 输入通道1（灰度图），输出通道32
self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
self.pool = nn.MaxPool2d(kernel_size=2, stride=2) # 下采样
self.fc1 = nn.Linear(64 * 7 * 7, 128) # 全连接层（输入维度7x7x64，输出128）
self.fc2 = nn.Linear(128, 10) # 输出层（10个类别）
def forward(self, x):
x = self.pool(F.relu(self.conv1(x))) # 第一层卷积+ReLU+池化
x = self.pool(F.relu(self.conv2(x))) # 第二层卷积+ReLU+池化
x = x.view(-1, 64 * 7 * 7) # 展平张量（-1表示自动计算批量大小）
x = F.relu(self.fc1(x)) # 全连接层+ReLU
x = self.fc2(x) # 输出层（无激活函数，用于分类）
return x
model = CNN() # 创建模型实例
print(model) # 查看模型结构
使用交叉熵损失（适用于分类任务）和Adam优化器（自适应学习率）。
import torch.optim as optim
criterion = nn.CrossEntropyLoss() # 交叉熵损失
optimizer = optim.Adam(model.parameters(), lr=0.001) # Adam优化器（学习率0.001）
# 训练循环（5个epoch）
for epoch in range(5):
model.train() # 设置为训练模式
running_loss = 0.0
for images, labels in train_loader:
optimizer.zero_grad() # 清空梯度
outputs = model(images) # 前向传播
loss = criterion(outputs, labels) # 计算损失
loss.backward() # 反向传播
optimizer.step() # 更新参数
running_loss += loss.item()
print(f"Epoch {epoch+1}, Loss: {running_loss/len(train_loader):.4f}") # 打印平均损失
model.eval() # 设置为评估模式（关闭Dropout等层）
correct = 0
total = 0
with torch.no_grad(): # 不计算梯度
for images, labels in test_loader:
outputs = model(images)
_, predicted = torch.max(outputs.data, 1) # 获取预测类别（argmax）
total += labels.size(0)
correct += (predicted == labels).sum().item()
print(f"Test Accuracy: {100 * correct / total:.2f}%") # 输出测试集准确率
# 保存模型参数
torch.save(model.state_dict(), 'mnist_cnn.pth')
# 加载模型参数
model = CNN() # 重新创建模型实例
model.load_state_dict(torch.load('mnist_cnn.pth'))
model.eval() # 设置为评估模式
nvidia-smi查看驱动支持的CUDA版本，选择对应的PyTorch安装命令（如CUDA 12.1对应cu121）。source pytorch_env/bin/activate激活环境，激活后终端提示符会显示环境名。pip install torch -i https://pypi.tuna.tsinghua.edu.cn/simple。通过以上步骤，你可在Ubuntu上完成PyTorch环境搭建，并实现一个完整的图像分类任务。后续可尝试更复杂的模型（如ResNet）或迁移学习（如使用预训练的BERT模型处理NLP任务）。
