CNN图像处理:使用ResNet-101模型对CIFAR-100数据集进行分类
使用ResNet-101模型进行CIFAR-100图像分类的代码详解
本代码示例使用预训练的ResNet-101模型对CIFAR-100数据集进行图像分类。代码包含数据预处理、模型训练和评估等步骤。
1. 导入必要的库
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.transforms as transforms
from torchvision.datasets import CIFAR100
from torch.utils.data import DataLoader
from torchvision.models import resnet101
import torchvision.models as models
- 导入PyTorch库,包括
torch、torch.nn、torch.optim、torchvision.transforms、torchvision.datasets、torch.utils.data、torchvision.models。
2. 设置参数
num_classes = 100
batch_size = 32
- 设置数据集的类别数(
num_classes)为100,表示CIFAR-100数据集包含100个类别。 - 设置批量大小(
batch_size)为32,表示每次训练使用32个样本。
3. 数据预处理
data_transforms = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.2023, 0.1994, 0.2010])
])
- 创建数据预处理管道
data_transforms,包含以下步骤:transforms.Resize(256):将图像调整为256x256大小。transforms.CenterCrop(224):从图像中心裁剪出224x224大小的区域。transforms.ToTensor():将图像转换为PyTorch张量。transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.2023, 0.1994, 0.2010]):对图像进行标准化,使用CIFAR-100数据集的均值和标准差进行归一化。
4. 加载数据集
train_dataset = CIFAR100(root='data/train_data', train=True, download=True, transform=data_transforms)
val_dataset = CIFAR100(root='data/val_data', train=False, download=True, transform=data_transforms)
- 使用
CIFAR100类加载CIFAR-100数据集,创建训练集(train_dataset)和验证集(val_dataset)的数据集对象。root参数指定数据集存储路径。train参数指定是否加载训练集。download参数指定是否下载数据集(如果不存在)。transform参数指定应用于图像的数据预处理管道。
5. 创建数据加载器
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=batch_size)
- 使用
DataLoader类创建数据加载器,用于批量加载训练集(train_loader)和验证集(val_loader)的数据。batch_size参数指定每次加载的样本数量。shuffle参数指定是否对训练集进行随机打乱。
6. 加载预训练模型
model = models.resnet101(weights=models.ResNet101_Weights.DEFAULT)
model.fc = nn.Linear(model.fc.in_features, num_classes)
- 使用
models.resnet101加载预训练的ResNet-101模型。 - 替换模型最后一层全连接层(
model.fc),使其输出维度与类别数(num_classes)一致。
7. 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
- 定义损失函数为交叉熵损失函数(
nn.CrossEntropyLoss)。 - 定义优化器为Adam优化器(
optim.Adam),学习率设置为0.001。
8. 将模型移动到GPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)
- 检测是否有可用的GPU,如果有则将模型移动到GPU上进行训练。
9. 开始训练
num_epochs = 10
for epoch in range(num_epochs):
model.train()
running_loss = 0.0
for batch_idx, (images, labels) in enumerate(train_loader):
images = images.to(device)
labels = labels.to(device)
optimizer.zero_grad()
outputs = model(images)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
print(f'Epoch {epoch+1}/{num_epochs}, Loss: {running_loss / len(train_loader):.4f}')
- 设置训练循环,迭代
num_epochs次。- 循环中,首先将模型设置为训练模式(
model.train())。 - 然后遍历训练数据集,计算模型输出、损失、反向传播和参数更新。
- 最后,打印每个epoch的平均损失。
- 循环中,首先将模型设置为训练模式(
10. 评估模型
model.eval()
correct = 0
total = 0
with torch.no_grad():
for batch_idx, (images, labels) in enumerate(val_loader):
images = images.to(device)
labels = labels.to(device)
outputs = model(images)
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
val_accuracy = 100 * correct / total
print(f'Validation Accuracy: {val_accuracy:.2f}%')
- 将模型设置为评估模式(
model.eval())。 - 遍历验证数据集,计算模型预测结果,并统计正确预测的样本数量。
- 最后,计算验证集上的准确率并打印。
11. 保存模型
torch.save(custom_model.state_dict(), 'custom_model1_home')
- 使用
torch.save函数将训练好的模型参数保存到文件。
总结
本代码示例使用预训练的ResNet-101模型对CIFAR-100数据集进行图像分类,并介绍了数据预处理、模型训练和评估等步骤。通过本示例可以了解使用预训练模型进行图像分类的基本步骤,并为进一步的图像分类任务提供参考。
原文地址: https://www.cveoy.top/t/topic/deYD 著作权归作者所有。请勿转载和采集!