首页 > 分享 > 使用pytorch实现基于VGG 19预训练模型的鲜花识别分类器，准确度达到97%

使用pytorch实现基于VGG 19预训练模型的鲜花识别分类器，准确度达到97%

花匠小妙招
2024-10-31 02:24

项目说明

本文使用的数据集是网络开源的鲜花数据集，并且基于VGG19的预训练模型通过迁移学习重新训练鲜花数据由此构建一个鲜花识别分类器

数据集

可以在此处找到有关花朵数据集的信息。数据集为102个花类的每一个都包含一个单独的文件夹。每朵花都标记为一个数字，每个编号的目录都包含许多.jpg文件。

实验环境

prtorch库
PIL库
如果想使用GPU训练的话请使用英伟达的显卡并安装好CUDA
如果用GPU的话我在自己电脑上使用GPU只使用了91分钟（我的GPU是1050）

##倒入库并检测是否有可用GPU

%matplotlib inline %config InlineBackend.figure_format = 'retina' import time import json import copy import matplotlib.pyplot as plt import seaborn as sns import numpy as np import PIL from PIL import Image from collections import OrderedDict import torch from torch import nn, optim from torch.optim import lr_scheduler from torch.autograd import Variable import torchvision from torchvision import datasets, models, transforms from torch.utils.data.sampler import SubsetRandomSampler import torch.nn as nn import torch.nn.functional as F import os # check if GPU is available train_on_gpu = torch.cuda.is_available() if not train_on_gpu: print('Bummer! Training on CPU ...') else: print('You are good to go! Training on GPU ...') ##有GPU就启用 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

123456789101112131415161718192021222324252627282930313233343536

定义数据集位置

data_dir = 'F:资料项目image_classifier_pytorch-master\flower_data' train_dir = 'flower_data/train' valid_dir = 'flower_data/valid' 123'

导入数据集并对数据进行处理

# Define your transforms for the training and testing sets 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]) ]) } # Load the datasets with ImageFolder image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x), data_transforms[x]) for x in ['train', 'valid']} # Using the image datasets and the trainforms, define the dataloaders batch_size = 64 dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=batch_size, shuffle=True, num_workers=4) for x in ['train', 'valid']} class_names = image_datasets['train'].classes dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'valid']} class_names = image_datasets['train'].classes # Label mapping with open('F:资料项目image_classifier_pytorch-mastercat_to_name.json', 'r') as f: cat_to_name = json.load(f)

1234567891011121314151617181920212223242526272829303132333435363738

查看数据情况

# Run this to test the data loader images, labels = next(iter(dataloaders['train'])) images.size() # # Run this to test your data loader images, labels = next(iter(dataloaders['train'])) rand_idx = np.random.randint(len(images)) # print(rand_idx) print("label: {}, class: {}, name: {}".format(labels[rand_idx].item(), class_names[labels[rand_idx].item()], cat_to_name[class_names[labels[rand_idx].item()]])) 12345678910

定义模型

model_name = 'densenet' #vgg if model_name == 'densenet': model = models.densenet161(pretrained=True) num_in_features = 2208 print(model) elif model_name == 'vgg': model = models.vgg19(pretrained=True) num_in_features = 25088 print(model.classifier) else: print("Unknown model, please choose 'densenet' or 'vgg'") # Create classifier for param in model.parameters(): param.requires_grad = False def build_classifier(num_in_features, hidden_layers, num_out_features): classifier = nn.Sequential() if hidden_layers == None: classifier.add_module('fc0', nn.Linear(num_in_features, 102)) else: layer_sizes = zip(hidden_layers[:-1], hidden_layers[1:]) classifier.add_module('fc0', nn.Linear(num_in_features, hidden_layers[0])) classifier.add_module('relu0', nn.ReLU()) classifier.add_module('drop0', nn.Dropout(.6)) classifier.add_module('relu1', nn.ReLU()) classifier.add_module('drop1', nn.Dropout(.5)) for i, (h1, h2) in enumerate(layer_sizes): classifier.add_module('fc'+str(i+1), nn.Linear(h1, h2)) classifier.add_module('relu'+str(i+1), nn.ReLU()) classifier.add_module('drop'+str(i+1), nn.Dropout(.5)) classifier.add_module('output', nn.Linear(hidden_layers[-1], num_out_features)) return classifier hidden_layers = None#[4096, 1024, 256][512, 256, 128] classifier = build_classifier(num_in_features, hidden_layers, 102) print(classifier) # Only train the classifier parameters, feature parameters are frozen if model_name == 'densenet': model.classifier = classifier criterion = nn.CrossEntropyLoss() optimizer = optim.Adadelta(model.parameters()) # Adadelta #weight optim.Adam(model.parameters(), lr=0.001, momentum=0.9) #optimizer_conv = optim.SGD(model.parameters(), lr=0.0001, weight_decay=0.001, momentum=0.9) sched = optim.lr_scheduler.StepLR(optimizer, step_size=4) elif model_name == 'vgg': model.classifier = classifier criterion = nn.NLLLoss() optimizer = optim.Adam(model.classifier.parameters(), lr=0.0001) sched = lr_scheduler.StepLR(optimizer, step_size=4, gamma=0.1) else: pass def train_model(model, criterion, optimizer, sched, num_epochs=5): since = time.time() best_model_wts = copy.deepcopy(model.state_dict()) best_acc = 0.0 for epoch in range(num_epochs): print('Epoch {}/{}'.format(epoch+1, num_epochs)) print('-' * 10) # Each epoch has a training and validation phase for phase in ['train', 'valid']: if phase == 'train': model.train() # Set model to training mode else: model.eval() # Set model to evaluate mode running_loss = 0.0 running_corrects = 0 # Iterate over data. for inputs, labels in dataloaders[phase]: inputs = inputs.to(device) labels = labels.to(device) # zero the parameter gradients optimizer.zero_grad() # forward # track history if only in train with torch.set_grad_enabled(phase == 'train'): outputs = model(inputs) _, preds = torch.max(outputs, 1) loss = criterion(outputs, labels) # backward + optimize only if in training phase if phase == 'train': #sched.step() loss.backward() optimizer.step() # statistics 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('{} Loss: {:.4f} Acc: {:.4f}'.format( phase, epoch_loss, epoch_acc)) # deep copy the model if phase == 'valid' and epoch_acc > best_acc: best_acc = epoch_acc best_model_wts = copy.deepcopy(model.state_dict()) print() time_elapsed = time.time() - since print('Training complete in {:.0f}m {:.0f}s'.format( time_elapsed // 60, time_elapsed % 60)) print('Best val Acc: {:4f}'.format(best_acc)) #load best model weights model.load_state_dict(best_model_wts) return model

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123