end0tknr's kipple - 新web写経開発

http://d.hatena.ne.jp/end0tknr/ から移転しました

独自データセットを CNN(AlexNet) + 転移学習 で画像分類

独自データセットを CNN(AlexNet) で画像分類 - end0tknr's kipple - 新web写経開発 の続きとして 前回エントリの内容を、「CNN(AlexNet) + 転移学習」で実施。

転移学習とは、学習済のモデルを再利用するもので、 今回の場合、CNN(AlexNet) の最終のみ、重みを更新する学習を行っています。

#!/usr/local/python3/bin/python3
# -*- coding: utf-8 -*-
import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler
from torch.autograd import Variable
import numpy as np
import torchvision
from torchvision import datasets, models, transforms

# import matplotlib.pyplot as plt
import matplotlib
matplotlib.use('Agg')
import matplotlib.pylab as plt

from PIL import Image
import time
import os

import cv2
from PIL import Image

device = 'cuda' if torch.cuda.is_available() else 'cpu'

num_epochs = 50


def main():
    print("DEVICE:", device)


    #画像の前処理を定義
    data_transforms = {
        'train': transforms.Compose([
            transforms.RandomResizedCrop(224),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ]),
        'val': transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ]),
    }


    #画像とラベルを読み込む
    data_dir = 'hymenoptera_data'
    image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x),
                                              data_transforms[x])
                      for x in ['train', 'val']}

    train_loader = torch.utils.data.DataLoader(image_datasets['train'], batch_size=5,
                                                 shuffle=True, num_workers=4)
    test_loader = torch.utils.data.DataLoader(image_datasets['val'], batch_size=5,
                                                 shuffle=False, num_workers=4)

    dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']}
    class_names = image_datasets['train'].classes


    net = models.alexnet(pretrained=True)  # 学習済み重みを利用するAlexNet
    net = net.to(device)

    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(net.parameters(), lr=0.01, momentum=0.9, weight_decay=5e-4)


    # ネットワークのパラメータを凍結
    for param in net.parameters():
        param.requires_grad = False  # backward時に重みを更新させない
    net = net.to(device)
    # 最終層を2クラス用(アリ,ハチ)に変更.
    # 変更することで、最終層はbackward時に重みが更新される
    num_ftrs = net.classifier[6].in_features
    net.classifier[6] = nn.Linear(num_ftrs, 2).to(device)


    # 最適化関数.
    # lr_schedulerは学習率を変更する為のものらしいが、理解していません
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(net.parameters(), lr=0.01, momentum=0.9, weight_decay=5e-4)
    lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)


    train_loss_list = []
    train_acc_list = []
    val_loss_list = []
    val_acc_list = []

    for epoch in range(num_epochs):
        
        train_loss = 0
        train_acc = 0
        val_loss = 0
        val_acc = 0

        #train
        net.train()
        for i, (images, labels) in enumerate(train_loader):
            images, labels = images.to(device), labels.to(device)

            optimizer.zero_grad()
            outputs = net(images)
            loss = criterion(outputs, labels)
            train_loss += loss.item()
            train_acc += (outputs.max(1)[1] == labels).sum().item()
            loss.backward()
            optimizer.step()

        avg_train_loss = train_loss / len(train_loader.dataset)
        avg_train_acc = train_acc / len(train_loader.dataset)

        #val
        net.eval()
        with torch.no_grad():
            for images, labels in test_loader:
                images = images.to(device)
                labels = labels.to(device)
                outputs = net(images)
                loss = criterion(outputs, labels)
                val_loss += loss.item()
                val_acc += (outputs.max(1)[1] == labels).sum().item()
        avg_val_loss = val_loss / len(test_loader.dataset)
        avg_val_acc = val_acc / len(test_loader.dataset)

        output_str_format = ', '.join(['Epoch [{}/{}]','Loss: {loss:.4f}',
                                       'val_loss: {val_loss:.4f}','val_acc: {val_acc:.4f}',
                                       'lr:{learning_rate}'])
        print(output_str_format.format(epoch+1, num_epochs, i+1, loss=avg_train_loss,
                                       val_loss=avg_val_loss, val_acc=avg_val_acc,
                                       learning_rate=optimizer.param_groups[0]["lr"]))
        
        #学習率調整
        lr_scheduler.step()
        train_loss_list.append(avg_train_loss)
        train_acc_list.append(avg_train_acc)
        val_loss_list.append(avg_val_loss)
        val_acc_list.append(avg_val_acc)
        
    output_to_file(train_loss_list,train_acc_list,val_loss_list,val_acc_list)

    
def output_to_file(train_loss_list,train_acc_list,val_loss_list,val_acc_list):

    plt.figure()
    plt.plot(range(num_epochs),
             train_loss_list,
             color='blue',
             linestyle='-',
             label='train_loss')
    plt.plot(range(num_epochs),
             val_loss_list,
             color='green',
             linestyle='--',
             label='val_loss')
#    plt.ylim([0,0.04])
    plt.legend()
    plt.xlabel('epoch')
    plt.ylabel('loss')
    plt.title('Training and validation loss')
    plt.grid()
    plt.savefig( 'test_4_1.png' )

    plt.figure()
    plt.plot(range(num_epochs),
             train_acc_list,
             color='blue',
             linestyle='-',
             label='train_acc')
    plt.plot(range(num_epochs),
             val_acc_list,
             color='green',
             linestyle='--',
             label='val_acc')
    plt.ylim([0,1])
    plt.legend()
    plt.xlabel('epoch')
    plt.ylabel('acc')
    plt.title('Training and validation accuracy')
    plt.grid()
    plt.savefig( 'test_4_2.png' )


if __name__ == '__main__':
    main()

↑こう書くと、↓こう出力されます。

DEVICE: cuda
Epoch [1/50], Loss: 0.8953, val_loss: 1.2442, val_acc: 0.8366, lr:0.01
  :
Epoch [20/50], Loss: 0.0634, val_loss: 1.4746, val_acc: 0.8954, lr:0.001
Epoch [21/50], Loss: 0.2592, val_loss: 1.4652, val_acc: 0.8954, lr:0.00010000000000000002
  :
Epoch [30/50], Loss: 0.2551, val_loss: 1.4208, val_acc: 0.8889, lr:0.00010000000000000002
Epoch [31/50], Loss: 0.2619, val_loss: 1.4198, val_acc: 0.8889, lr:1.0000000000000003e-05
  :
Epoch [50/50], Loss: 0.2745, val_loss: 1.4169, val_acc: 0.8889, lr:1.0000000000000002e-06

f:id:end0tknr:20191020075721p:plainf:id:end0tknr:20191020075724p:plain