PyTorch图像分类实战：使用预训练ResNet-101模型在CIFAR-100数据集上训练

这段代码包含了使用PyTorch进行图像分类的实例。代码的主要功能是使用预训练的ResNet-101模型在CIFAR-100数据集上进行训练和验证。

以下是代码的主要功能和步骤：

导入所需的PyTorch模块和库。
设置参数，例如类别数和批量大小。
定义数据预处理步骤，包括图像的调整大小、裁剪和归一化。
加载CIFAR-100数据集，并创建相应的数据加载器。
加载预训练的ResNet-101模型，并根据数据集的类别数替换最后一层全连接层。
定义损失函数和优化器。
将模型移动到GPU（如果可用）。
进行训练循环，遍历每个epoch和每个batch，计算损失并更新模型参数。
在训练结束后，使用验证集计算模型的准确率。
调整学习率。
保存训练好的模型。

需要注意的是，代码中的路径和文件名可能需要根据你的实际情况进行修改。此外，还可以根据需求进行额外的修改和优化。

如果你遇到任何错误或问题，请提供具体的错误信息或问题描述，以便我能够更具体地帮助你。

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.transforms as transforms
from torchvision.datasets import ImageNet
from torch.utils.data import DataLoader
from torchvision.models import resnet101
from torchvision.datasets import CIFAR100
import torchvision.models as models

# 设置参数
#num_classes = 1000  # ImageNet数据集的类别数
num_classes = 100  # CIFAR-100数据集的类别数

batch_size = 32

# 数据预处理
data_transforms = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    #transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.2023, 0.1994, 0.2010])
])


# 加载ImageNet数据集
#train_dataset = ImageNet(root='path/to/train_data', split='train', transform=data_transforms)
#val_dataset = ImageNet(root='path/to/val_data', split='val', transform=data_transforms)
# 加载CIFAR-100数据集
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)


# 创建数据加载器
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=batch_size)



# 加载预训练的ResNet-101模型
#model = resnet101(pretrained=True)

# 加载预训练的ResNet-101模型
#model = resnet101(weights='imagenet')
model = models.resnet101(weights = models.ResNet101_Weights.DEFAULT)
# 将最后一层全连接层替换为新的分类器
model.fc = nn.Linear(model.fc.in_features, num_classes)

# 定义损失函数
criterion = nn.CrossEntropyLoss()

custom_model = model

# 定义优化器
#optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
optimizer = optim.Adam(model.parameters(), lr=0.001)


# 定义学习率调度器
lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1)


# 训练模型
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

print("当前使用" + str(device) + "进行训练")

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()
        epoch_loss = running_loss / len(train_loader)
    print(f"Epoch {epoch+1}/{num_epochs}, Loss: {epoch_loss:.4f}")


model.eval()
correct = 0
total = 0
with torch.no_grad():
    #for images, labels in val_loader:
    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}%")

lr_scheduler.step()
# 假设训练好的模型为model
# 假设保存路径为checkpoint.pth

# 保存模型
torch.save(custom_model.state_dict(), 'custom_model1_home')

PyTorch图像分类实战：使用预训练ResNet-101模型在CIFAR-100数据集上训练