首页
学习
活动
专区
圈层
工具
发布
首页
学习
活动
专区
圈层
工具
MCP广场
社区首页 >专栏 >深度学习--使用PyTorch训练模型

深度学习--使用PyTorch训练模型

原创
作者头像
china马斯克
发布2025-03-31 11:45:22
发布2025-03-31 11:45:22
3950
举报

友友们,周一好呀!又见面了。今天我们来继续充电。这篇讲一下我们如何来利用PyTorch训练图像识别的模型。ok,下面进入正文。

先来看一下我们训练的基本步骤吧。

一、训练神经网络的基本步骤

1 数据准备:首先,我们需要准备好训练数据和验证数据。这些数据通常会被加载到DataLoader中,以便于批量处理和迭代。

2 模型定义:接下来,我们需要定义一个神经网络模型。在PyTorch中,这通常是通过继承torch.nn.Module类并实现__init__和forward方法来完成的。

3 损失函数与优化器:一般来说选择合适的损失函数和优化器是训练神经网络的关键。损失函数用于衡量模型预测与真实标签之间的差异,而优化器则负责根据梯度更新模型参数。

4 训练循环:训练循环包括前向传播、计算损失、反向传播和参数更新四个步骤。这个过程通常会迭代多次(即epoch),直到模型在验证集上的性能达到满意为止。

下面我们就按照上面的步骤来实战一下吧。

二、实战

首先是准备数据,这里我们使用CIFAR-10数据集,它是一个常用的用于图像分类的基准数据集。

先来导入库:

代码语言:txt
复制
import torch
import torch.nn as nn
import torch.optim as optim
import torch.utils.data as Data
from torchvision import datasets, transforms

然后我们加载数据:

代码语言:txt
复制
# 数据预处理
transform = transforms.Compose([
    transforms.Resize((32, 32)), 
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# 加载训练集和测试集
train_dataset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
test_dataset = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
# 使用DataLoader加载数据
train_loader = Data.DataLoader(dataset=train_dataset, batch_size=128, shuffle=True, num_workers=4)
test_loader = Data.DataLoader(dataset=test_dataset, batch_size=128, shuffle=False, num_workers=4)

接着进行构建:

代码语言:txt
复制
class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
        self.fc1 = nn.Linear(64 * 8 * 8, 512)
        self.fc2 = nn.Linear(512, 10)  # CIFAR-10有10个类别

    def forward(self, x):
        x = self.pool(torch.relu(self.conv1(x)))
        x = self.pool(torch.relu(self.conv2(x)))
        x = x.view(-1, 64 * 8 * 8)  # 展平为一维向量
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# 实例化模型
model = SimpleCNN().to(torch.device('cuda' if torch.cuda.is_available() else 'cpu'))

定义损失函数与优化器:

代码语言:txt
复制
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

在PyTorch中,nn.CrossEntropyLoss()期望的输入是原始的对数几率(logits),而不是softmax归一化后的概率。该损失函数内部会计算softmax。

代码中model.parameters()返回模型中所有可训练参数的迭代器,这些参数将被优化器更新。

lr=0.001设置了学习率的初始值为0.001。学习率是一个超参数,它决定了参数更新的步长大小。较大的学习率可能导致训练过程不稳定,而较小的学习率可能导致训练过程缓慢。

ok,下面就可以开始训练模型了:

代码语言:txt
复制
num_epochs = 20

for epoch in range(num_epochs):
    model.train()
    running_loss = 0.0
    for i, (inputs, labels) in enumerate(train_loader):
        inputs, labels = inputs.to(torch.device('cuda' if torch.cuda.is_available() else 'cpu')), labels.to(torch.device('cuda' if torch.cuda.is_available() else 'cpu'))
        
        optimizer.zero_grad()
        
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        
        running_loss += loss.item()
        if i % 100 == 99:  # 每100个batch打印一次
            print(f'[Epoch {epoch + 1}, Batch {i + 1}] loss: {running_loss / 100:.3f}')
            running_loss = 0.0

print('Finished Training')

在完成训练后,我们来测试一下上面的模型训练效果如何:

代码语言:txt
复制
model.eval()
correct = 0
total = 0
with torch.no_grad():
    for inputs, labels in test_loader:
        inputs, labels = inputs.to(torch.device('cuda' if torch.cuda.is_available() else 'cpu')), labels.to(torch.device('cuda' if torch.cuda.is_available() else 'cpu'))
        outputs = model(inputs)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print(f'Accuracy of the model on the test images: {100 * correct / total} %')

当然,要注意,我们在训练过程中,要记录详细的日志信息(如每个epoch的损失值、准确率等)对于监控模型训练进度和调试问题非常有帮助。这里我的日志信息是通过简单的打印语句输出的。

如果效果不错,我们就可以保存模型了:

代码语言:txt
复制
torch.save(model.state_dict(), 'mytest.pth')

这就是一个完整的使用PyTorch训练模型的流程了,从导入库,数据加载到模型训练、测试和保存。希望对大家有所帮助。

原创声明:本文系作者授权腾讯云开发者社区发表,未经许可,不得转载。

如有侵权,请联系 cloudcommunity@tencent.com 删除。

原创声明:本文系作者授权腾讯云开发者社区发表,未经许可,不得转载。

如有侵权,请联系 cloudcommunity@tencent.com 删除。

评论
登录后参与评论
0 条评论
热度
最新
推荐阅读
目录
  • 一、训练神经网络的基本步骤
  • 二、实战
领券
问题归档专栏文章快讯文章归档关键词归档开发者手册归档开发者手册 Section 归档