友友们,周一好呀!又见面了。今天我们来继续充电。这篇讲一下我们如何来利用PyTorch训练图像识别的模型。ok,下面进入正文。
先来看一下我们训练的基本步骤吧。
1 数据准备:首先,我们需要准备好训练数据和验证数据。这些数据通常会被加载到DataLoader中,以便于批量处理和迭代。
2 模型定义:接下来,我们需要定义一个神经网络模型。在PyTorch中,这通常是通过继承torch.nn.Module类并实现__init__和forward方法来完成的。
3 损失函数与优化器:一般来说选择合适的损失函数和优化器是训练神经网络的关键。损失函数用于衡量模型预测与真实标签之间的差异,而优化器则负责根据梯度更新模型参数。
4 训练循环:训练循环包括前向传播、计算损失、反向传播和参数更新四个步骤。这个过程通常会迭代多次(即epoch),直到模型在验证集上的性能达到满意为止。
下面我们就按照上面的步骤来实战一下吧。
首先是准备数据,这里我们使用CIFAR-10数据集,它是一个常用的用于图像分类的基准数据集。
先来导入库:
import torch
import torch.nn as nn
import torch.optim as optim
import torch.utils.data as Data
from torchvision import datasets, transforms
然后我们加载数据:
# 数据预处理
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)
接着进行构建:
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'))
定义损失函数与优化器:
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,下面就可以开始训练模型了:
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')
在完成训练后,我们来测试一下上面的模型训练效果如何:
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的损失值、准确率等)对于监控模型训练进度和调试问题非常有帮助。这里我的日志信息是通过简单的打印语句输出的。
如果效果不错,我们就可以保存模型了:
torch.save(model.state_dict(), 'mytest.pth')
这就是一个完整的使用PyTorch训练模型的流程了,从导入库,数据加载到模型训练、测试和保存。希望对大家有所帮助。
原创声明:本文系作者授权腾讯云开发者社区发表,未经许可,不得转载。
如有侵权,请联系 cloudcommunity@tencent.com 删除。
原创声明:本文系作者授权腾讯云开发者社区发表,未经许可,不得转载。
如有侵权,请联系 cloudcommunity@tencent.com 删除。