独自データセットを 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
