使用pytorch實現基於VGG 19預訓練模型的鮮花識別分類器

項目說明

本文使用的數據集是網絡開源的鮮花數據集，而且基於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")

定義數據集位置

data_dir = 'F:\資料\項目\image_classifier_pytorch-master\\flower_data'

train_dir = 'flower_data/train'

valid_dir = 'flower_data/valid'

導入數據集並對數據進行處理

# 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-master\cat_to_name.json', 'r') as f:

    cat_to_name = json.load(f)

查看數據狀況

# 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()]]))

定義模型

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())

function(){ //XM返傭 http://www.kaifx.cn/broker/xm.html

        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

開始訓練

epochs = 30

model.to(device)

model = train_model(model, criterion, optimizer, sched, epochs)