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训练卷积神经网络的完整实现指南和使用例子。您可以根据自己的需求进行修改和调整,以获得更好的结果。祝您好运!
