欢迎访问宙启技术站
智能推送

CIFAR100数据集上的卷积神经网络训练:PyTorch实现指南

发布时间:2023-12-29 13:02:07

CIFAR-100数据集是一个常用的图像分类数据集,包含100个类别的60000个32x32彩色图像。在本指南中,我将为您提供使用PyTorch对CIFAR-100数据集进行训练的实现指南,并带有一个完整的使用例子。

首先,我们需要安装PyTorch,并确保版本为1.8.0或更高。您可以使用以下命令安装PyTorch:

pip install torch torchvision

接下来,让我们导入所需的库和模块:

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.transforms as transforms
import torchvision.datasets as datasets

在我们开始之前,我们需要定义一些超参数,例如批量大小、学习率、训练 epoch 数等:

batch_size = 128
learning_rate = 0.001
num_epochs = 10

然后,我们需要创建一个数据转换对象,用于对CIFAR-100数据集进行预处理。我们将通过转换将图像数据转换为张量,并将图像像素值归一化到0到1之间:

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

接下来,我们将下载和加载CIFAR-100训练和测试数据集:

train_dataset = datasets.CIFAR100(root='./data', train=True, transform=transform, download=True)
test_dataset = datasets.CIFAR100(root='./data', train=False, transform=transform, download=True)

现在,我们将创建训练和测试数据加载器,用于加载数据集的批次数据:

train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False)

现在,我们定义一个简单的卷积神经网络模型,并将其命名为CIFARNet。在这个例子中,我们使用了三个卷积层、ReLU激活函数、最大池化层和全连接层。网络结构如下:

class CIFARNet(nn.Module):
    def __init__(self):
        super(CIFARNet, self).__init__()
        self.conv1 = nn.Conv2d(3, 16, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=1)
        self.conv3 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.maxpool = nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc = nn.Linear(64 * 4 * 4, 100)
        
    def forward(self, x):
        x = self.conv1(x)
        x = self.maxpool(x)
        x = nn.ReLU()(x)
        x = self.conv2(x)
        x = self.maxpool(x)
        x = nn.ReLU()(x)
        x = self.conv3(x)
        x = self.maxpool(x)
        x = nn.ReLU()(x)
        x = self.fc(x.view(x.size(0), -1))
        return x

接下来,我们实例化模型,并定义损失函数和优化器:

model = CIFARNet()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

训练和测试循环如下:

for epoch in range(num_epochs):
    model.train()
    for batch_idx, (data, targets) in enumerate(train_loader):
        optimizer.zero_grad()
        outputs = model(data)
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()
        
    model.eval()
    with torch.no_grad():
        correct, total = 0, 0
        for batch_idx, (data, targets) in enumerate(test_loader):
            outputs = model(data)
            _, predicted = torch.max(outputs.data, 1)
            total += targets.size(0)
            correct += (predicted == targets).sum().item()
        
    accuracy = 100 * correct / total
    print(f'Epoch [{epoch+1}/{num_epochs}], Accuracy: {accuracy:.2f}%')

最后,在训练结束后,您可以保存模型的权重到磁盘上:

torch.save(model.state_dict(), 'cifar_net.pth')

这就是在CIFAR-100数据集上使用PyTorch训练卷积神经网络的完整实现指南和使用例子。您可以根据自己的需求进行修改和调整,以获得更好的结果。祝您好运!