如何在Ubuntu上使用PyTorch进行机器学习实验 - 问答

1. 安装PyTorch及依赖

1.1 系统准备

首先更新系统包列表并安装基础依赖，确保后续安装顺利：

sudo apt update
sudo apt install -y build-essential cmake git wget unzip yasm pkg-config libopenblas-dev liblapack-dev libjpeg-dev libpng-dev

1.2 安装Python与虚拟环境

Ubuntu通常预装Python3，需确认pip是否安装，随后创建虚拟环境隔离项目依赖：

sudo apt install -y python3 python3-pip
python3 -m venv pytorch_env  # 创建虚拟环境
source pytorch_env/bin/activate  # 激活环境（激活后命令行前缀变为“(pytorch_env)”）

1.3 安装PyTorch

根据是否有GPU选择安装方式：

CPU版本（无GPU加速）：

pip3 install torch torchvision torchaudio

GPU版本（需CUDA支持）：访问PyTorch官网获取对应CUDA版本的安装命令（如CUDA 11.7）：
```
pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu117
```

1.4 验证安装

运行Python解释器，检查PyTorch版本及GPU可用性：

import torch
print(torch.__version__)  # 输出版本号（如2.1.0）
print(torch.cuda.is_available())  # 若为True，说明GPU可用

2. 准备机器学习实验数据

以经典的MNIST手写数字识别为例，使用torchvision加载并预处理数据：

import torch
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

# 定义数据预处理：转Tensor并归一化（均值为0.5，标准差为0.5）
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

# 加载训练集（自动下载）和测试集
train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)

# 创建数据加载器（批量大小64，训练集打乱顺序）
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=1000, shuffle=False)

3. 定义机器学习模型

以CNN（卷积神经网络）为例，用于提取MNIST图像的空间特征：

import torch.nn as nn
import torch.nn.functional as F

class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        # 第一层卷积：输入1通道（灰度图），输出32通道，卷积核3x3
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1)
        # 第二层卷积：输入32通道，输出64通道，卷积核3x3
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
        # 最大池化层：窗口2x2，步幅2
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        # 全连接层1：输入1600（64*5*5），输出128
        self.fc1 = nn.Linear(64 * 5 * 5, 128)
        # 全连接层2：输入128，输出10（对应10个数字类别）
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        # 卷积->激活->池化
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        # 展平：将4D张量转为2D（batch_size, 64*5*5）
        x = x.view(-1, 64 * 5 * 5)
        # 全连接层->激活
        x = F.relu(self.fc1(x))
        # 输出层（无需激活，后续用CrossEntropyLoss自带Softmax）
        x = self.fc2(x)
        return x

# 实例化模型
model = CNN()

4. 训练模型

定义损失函数、优化器，编写训练循环迭代数据：

import torch.optim as optim

# 定义损失函数（交叉熵损失，适用于分类任务）
criterion = nn.CrossEntropyLoss()
# 定义优化器（SGD，学习率0.01，动量0.5）
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.5)

# 训练循环（5个epoch）
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")  # 自动选择设备
model.to(device)  # 将模型移至设备

for epoch in range(5):  # 遍历数据集5次
    model.train()  # 切换至训练模式
    running_loss = 0.0
    correct = 0
    total = 0
    
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)  # 将数据移至设备
        
        # 梯度清零
        optimizer.zero_grad()
        
        # 前向传播
        outputs = model(data)
        
        # 计算损失
        loss = criterion(outputs, target)
        
        # 反向传播+优化
        loss.backward()
        optimizer.step()
        
        # 统计指标
        running_loss += loss.item()
        _, predicted = outputs.max(1)  # 获取预测类别
        total += target.size(0)
        correct += predicted.eq(target).sum().item()
        
        # 打印批次信息（每100个批次）
        if batch_idx % 100 == 99:
            print(f'Epoch: {epoch+1}, Batch: {batch_idx+1}, Loss: {running_loss/100:.3f}, Acc: {100.*correct/total:.3f}%')
            running_loss = 0.0

5. 评估模型

在测试集上验证模型泛化能力：

model.eval()  # 切换至评估模式（关闭Dropout等层）
eval_loss = 0
correct = 0
total = 0

with torch.no_grad():  # 禁用梯度计算
    for data, target in test_loader:
        data, target = data.to(device), target.to(device)
        outputs = model(data)
        loss = criterion(outputs, target)
        
        eval_loss += loss.item()
        _, predicted = outputs.max(1)
        total += target.size(0)
        correct += predicted.eq(target).sum().item()

print(f'Test Loss: {eval_loss/len(test_loader):.4f}, Test Acc: {100.*correct/total:.2f}%')

6. 保存与加载模型

训练完成后，保存模型参数以便后续使用：

# 保存模型参数（仅保存state_dict，而非整个模型）
torch.save(model.state_dict(), 'mnist_cnn.pth')

# 加载模型（需重新实例化模型结构）
model = CNN()
model.load_state_dict(torch.load('mnist_cnn.pth'))
model.to(device)  # 移至设备

7. 可选优化与扩展

调整超参数：修改学习率（如lr=0.001）、批量大小（如batch_size=128）、网络层数（如增加conv3）。
使用回调函数：通过torch.utils.tensorboard记录训练过程，可视化损失和准确率。

数据增强：在transform中添加旋转、平移等操作，提升模型鲁棒性：

transform = transforms.Compose([
    transforms.RandomRotation(10),  # 随机旋转±10度
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

通过以上步骤，可在Ubuntu上完成PyTorch的机器学习实验流程，覆盖从环境搭建到模型部署的全环节。

0 赞

0 踩