从零开始:用PyTorch和VGG19构建你的第一个鲜花分类模型(附完整代码)

发布时间:2026/7/7 8:10:54

从零开始:用PyTorch和VGG19构建你的第一个鲜花分类模型(附完整代码) 从零开始用PyTorch和VGG19构建你的第一个鲜花分类模型附完整代码当你第一次看到那些绚丽的鲜花照片时是否曾想过让计算机也能像人类一样识别它们深度学习让这一切成为可能。本文将带你从零开始使用PyTorch框架和预训练的VGG19模型构建一个能够识别102种不同花卉的分类系统。无论你是刚接触深度学习的新手还是希望扩展实战经验的中级开发者这个项目都能让你获得宝贵的实践经验。1. 项目准备与环境搭建在开始编码之前我们需要确保开发环境配置正确。以下是推荐的环境配置Python 3.7建议使用Anaconda管理Python环境PyTorch 1.8深度学习框架核心CUDA 11.1可选如果使用NVIDIA GPU加速训练torchvision提供图像预处理工具和预训练模型安装PyTorch最简单的方式是通过官方命令pip install torch torchvision对于GPU加速版本需要根据CUDA版本选择对应的PyTorch安装命令。可以使用以下代码检查CUDA是否可用import torch print(torch.cuda.is_available()) # 输出True表示CUDA可用提示如果没有高性能GPU可以考虑使用Google Colab的免费GPU资源进行模型训练。2. 数据集获取与预处理我们将使用牛津大学的102 Category Flower Dataset这是花卉分类任务的基准数据集之一包含102类英国常见花卉每类有40到258张图像。2.1 数据集结构数据集通常分为三个子集训练集train用于模型训练验证集valid用于调整超参数测试集test用于最终评估模型性能典型的数据目录结构如下data/ train/ class1/ image1.jpg image2.jpg ... class2/ ... valid/ class1/ ... test/ class1/ ...2.2 图像预处理不同的数据集需要不同的预处理方式。对于训练集我们通常应用数据增强技术来增加样本多样性而对于验证集和测试集只需进行简单的归一化处理。from torchvision import transforms data_transforms { train: transforms.Compose([ transforms.RandomRotation(30), transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]), valid: transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]), test: transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) }注意归一化使用的均值和标准差是基于ImageNet数据集计算的这在迁移学习中很常见。3. 构建VGG19模型VGG19是牛津大学视觉几何组开发的经典卷积神经网络在ImageNet竞赛中表现出色。我们将使用预训练的VGG19作为基础模型并针对花卉分类任务进行微调。3.1 加载预训练模型import torchvision.models as models # 加载预训练的VGG19模型 vgg19 models.vgg19(pretrainedTrue) # 冻结所有卷积层的参数 for param in vgg19.features.parameters(): param.requires_grad False3.2 修改分类器原始的VGG19是为1000类ImageNet设计的我们需要修改最后的全连接层以适应我们的102类花卉分类任务。import torch.nn as nn # 修改分类器部分 vgg19.classifier[6] nn.Linear(4096, 102) # 将模型转移到GPU如果可用 device torch.device(cuda:0 if torch.cuda.is_available() else cpu) model vgg19.to(device)3.3 定义损失函数和优化器import torch.optim as optim criterion nn.CrossEntropyLoss() optimizer optim.SGD(model.parameters(), lr0.001, momentum0.9)4. 训练与评估模型4.1 训练过程实现训练深度学习模型需要迭代多个epoch每个epoch都会遍历整个训练集。我们还需要在验证集上定期评估模型性能。def train_model(model, criterion, optimizer, num_epochs25): best_acc 0.0 for epoch in range(num_epochs): print(fEpoch {epoch1}/{num_epochs}) print(- * 10) # 每个epoch都有训练和验证阶段 for phase in [train, valid]: if phase train: model.train() # 设置模型为训练模式 else: model.eval() # 设置模型为评估模式 running_loss 0.0 running_corrects 0 # 迭代数据 for inputs, labels in dataloaders[phase]: inputs inputs.to(device) labels labels.to(device) # 梯度清零 optimizer.zero_grad() # 前向传播 with torch.set_grad_enabled(phase train): outputs model(inputs) _, preds torch.max(outputs, 1) loss criterion(outputs, labels) # 只在训练阶段反向传播和优化 if phase train: loss.backward() optimizer.step() # 统计 running_loss loss.item() * inputs.size(0) running_corrects torch.sum(preds labels.data) epoch_loss running_loss / dataset_sizes[phase] epoch_acc running_corrects.double() / dataset_sizes[phase] print(f{phase} Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}) # 深度复制模型 if phase valid and epoch_acc best_acc: best_acc epoch_acc best_model_wts copy.deepcopy(model.state_dict()) print() print(fBest val Acc: {best_acc:4f}) # 加载最佳模型权重 model.load_state_dict(best_model_wts) return model4.2 模型评估训练完成后我们需要在独立的测试集上评估模型性能这是对模型泛化能力的真实检验。def evaluate_model(model, criterion, phasetest): model.eval() # 设置模型为评估模式 running_loss 0.0 running_corrects 0 # 迭代数据 for inputs, labels in dataloaders[phase]: inputs inputs.to(device) labels labels.to(device) # 前向传播 with torch.no_grad(): outputs model(inputs) _, preds torch.max(outputs, 1) loss criterion(outputs, labels) # 统计 running_loss loss.item() * inputs.size(0) running_corrects torch.sum(preds labels.data) epoch_loss running_loss / dataset_sizes[phase] epoch_acc running_corrects.double() / dataset_sizes[phase] print(f{phase} Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}) return epoch_acc4.3 可视化预测结果理解模型在哪些类别上表现良好哪些类别容易混淆对于改进模型非常重要。def visualize_model(model, num_images6): was_training model.training model.eval() images_so_far 0 fig plt.figure(figsize(15, 10)) with torch.no_grad(): for i, (inputs, labels) in enumerate(dataloaders[test]): inputs inputs.to(device) labels labels.to(device) outputs model(inputs) _, preds torch.max(outputs, 1) for j in range(inputs.size()[0]): images_so_far 1 ax plt.subplot(num_images//2, 2, images_so_far) ax.axis(off) ax.set_title(fpredicted: {class_names[preds[j]]}\nactual: {class_names[labels[j]]}) imshow(inputs.cpu().data[j]) if images_so_far num_images: model.train(modewas_training) return model.train(modewas_training)5. 模型优化与调参技巧获得初步模型后我们可以通过多种方式进一步提升性能5.1 学习率调整学习率是训练神经网络最重要的超参数之一。我们可以使用学习率调度器动态调整学习率from torch.optim import lr_scheduler # 每7个epoch将学习率乘以0.1 exp_lr_scheduler lr_scheduler.StepLR(optimizer, step_size7, gamma0.1)5.2 数据增强扩展增加更多数据增强技术可以提高模型泛化能力train_transforms transforms.Compose([ transforms.RandomRotation(30), transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), transforms.ColorJitter(brightness0.2, contrast0.2, saturation0.2, hue0.1), transforms.RandomAffine(degrees0, translate(0.1, 0.1)), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ])5.3 模型集成结合多个模型的预测结果往往能获得更好的性能from torch import nn class EnsembleModel(nn.Module): def __init__(self, modelA, modelB): super(EnsembleModel, self).__init__() self.modelA modelA self.modelB modelB self.classifier nn.Linear(204, 102) # 假设每个模型输出102维 def forward(self, x): x1 self.modelA(x) x2 self.modelB(x) x torch.cat((x1, x2), dim1) x self.classifier(x) return x6. 实际应用与模型部署训练好的模型可以应用于实际场景如花卉识别APP或智能园艺系统。6.1 保存和加载模型# 保存整个模型 torch.save(model, flower_classifier.pth) # 只保存模型参数推荐方式 torch.save(model.state_dict(), flower_classifier_params.pth) # 加载模型 model torch.load(flower_classifier.pth) # 或 model.load_state_dict(torch.load(flower_classifier_params.pth))6.2 创建预测函数def predict_image(image_path, model, topk5): 预测图像类别及概率 image Image.open(image_path) image test_transforms(image).unsqueeze(0) image image.to(device) model.eval() with torch.no_grad(): output model(image) probabilities torch.nn.functional.softmax(output[0], dim0) top_probs, top_classes probabilities.topk(topk) return top_probs.cpu().numpy(), top_classes.cpu().numpy()6.3 转换为ONNX格式将PyTorch模型转换为ONNX格式可以方便地在不同平台上部署# 创建虚拟输入 dummy_input torch.randn(1, 3, 224, 224, devicedevice) # 导出模型 torch.onnx.export(model, dummy_input, flower_classifier.onnx, input_names[input], output_names[output], dynamic_axes{input: {0: batch_size}, output: {0: batch_size}})

相关新闻