PyTorch深度学习框架进阶——图像分类模型

喜欢花科技君 2025-03-08 02:21:58

图像分类通常使用卷积神经网络(CNN)。设计一个简单的CNN结构,比如几个卷积层,然后是全连接层。或者使用预训练模型,比如ResNet。

数据预处理和加载:PyTorch提供了torchvision库,里面有transforms,可以用来做数据增强和标准化。比如,对训练集做随机裁剪、水平翻转,然后转换成Tensor,并归一化。

安装必要的库

确保已经安装了torch和torchvision。如果没有安装,可以使用以下命令进行安装:

pip install torch torchvision导入库并设置GPU设备import torchimport torch.nn as nnimport torch.optim as optimimport torchvisionimport torchvision.transforms as transformsfrom torchvision import modelsimport timeimport os# 检查是否有可用的GPUdevice = torch.device("cuda" if torch.cuda.is_available() else "cpu")print(f"使用设备: {device}")加载和预处理数据

使用CIFAR-10数据集,并进行一些数据增强和标准化。

# 数据预处理和增强transform_train = transforms.Compose([ transforms.RandomHorizontalFlip(), # 随机水平翻转 transforms.RandomCrop(32, padding=4), # 随机裁剪 transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),])transform_test = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),])# 下载并加载训练数据trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform_train)trainloader = torch.utils.data.DataLoader(trainset, batch_size=128, shuffle=True, num_workers=2)# 下载并加载测试数据testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform_test)testloader = torch.utils.data.DataLoader(testset, batch_size=100, shuffle=False, num_workers=2)# CIFAR-10的类别classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')定义模型

使用预训练的ResNet18模型,并替换其最后的全连接层以适应CIFAR-10的10个类别。

# 加载预训练的ResNet18模型model = models.resnet18(pretrained=True)# 冻结所有参数for param in model.parameters(): param.requires_grad = False# 替换最后的全连接层num_ftrs = model.fc.in_featuresmodel.fc = nn.Linear(num_ftrs, 10) # CIFAR-10有10个类别model = model.to(device)设置损失函数和优化器

由于我们冻结了预训练模型的参数,只有最后的全连接层的参数需要更新,因此优化器的参数只包含model.fc.parameters()。

criterion = nn.CrossEntropyLoss()optimizer = optim.Adam(model.fc.parameters(), lr=0.001)训练模型num_epochs = 10since = time.time()for epoch in range(num_epochs): model.train() running_loss = 0.0 running_corrects = 0 for inputs, labels in trainloader: inputs = inputs.to(device) labels = labels.to(device) optimizer.zero_grad() # 前向传播 outputs = model(inputs) _, preds = torch.max(outputs, 1) loss = criterion(outputs, labels) # 反向传播和优化 loss.backward() optimizer.step() running_loss += loss.item() * inputs.size(0) running_corrects += torch.sum(preds == labels.data) epoch_loss = running_loss / len(trainset) epoch_acc = running_corrects.double() / len(trainset) print(f'Epoch {epoch+1}/{num_epochs} 损失: {epoch_loss:.4f} 准确率: {epoch_acc:.4f}') # 在每个epoch结束时在测试集上评估 model.eval() test_loss = 0.0 test_corrects = 0 with torch.no_grad(): for inputs, labels in testloader: inputs = inputs.to(device) labels = labels.to(device) outputs = model(inputs) _, preds = torch.max(outputs, 1) loss = criterion(outputs, labels) test_loss += loss.item() * inputs.size(0) test_corrects += torch.sum(preds == labels.data) test_epoch_loss = test_loss / len(testset) test_epoch_acc = test_corrects.double() / len(testset) print(f'测试集 损失: {test_epoch_loss:.4f} 准确率: {test_epoch_acc:.4f}')time_elapsed = time.time() - sinceprint(f'训练完成,总耗时: {time_elapsed//60:.0f}m {time_elapsed%60:.0f}s')评估模型

训练完成后,在测试集上评估模型的性能。

# 最终在测试集上的评估model.eval()correct = 0total = 0with torch.no_grad(): for inputs, labels in testloader: inputs = inputs.to(device) labels = labels.to(device) outputs = model(inputs) _, predicted = torch.max(outputs.data, 1) total += labels.size(0) correct += (predicted == labels).sum().item()print(f'测试集上的准确率: {100 * correct / total:.2f}%')

0 阅读:2

喜欢花科技君

简介:感谢大家的关注