【ネコと機械学習2】Pytorchを使ったネコ判別器の作り方

#必要なライブラリ用意
import numpy as np
import random
import glob
from PIL import Image
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torch.utils.data as data
import torchvision
from tqdm import tqdm
from torchvision import models, transforms, datasets
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report
import os
import time

#シードを固定

torch.manual_seed(0)
np.random.seed(0)
random.seed(0)

#データの前処理を行うクラスを定義

class ImageTransform():
    def __init__(self, resize, mean, std):
        self.data_transform = {
            'train': transforms.Compose([
                transforms.Resize((resize, resize)),
                transforms.RandomHorizontalFlip(p=0.5),
                transforms.RandomCrop(size =(150, 150), padding=18),
                transforms.ToTensor(),
                #transforms.Normalize(mean,std)
            ]),
            'val': transforms.Compose([
                transforms.Resize((resize,resize)),
                transforms.ToTensor(),
                #transforms.transforms.Normalize(mean, std)
            ]),
            'test' : transforms.Compose([
                transforms.Resize((resize,resize)),
                transforms.ToTensor(),
                #transforms.transforms.Normalize(mean, std)
            ])
        }

    def __call__(self, img, phase):
        return self.data_transform[phase](img)

#データにネコ=0、イヌ=1のラベル付けを行うクラスを定義

class DogCatDataset(data.Dataset):
    def __init__(self, file_list, phase, transform = None):
        self.file_list = file_list
        self.transform = transform
        self.phase = phase

    def __len__(self):
        return len(self.file_list)

    def __getitem__(self, index):
        img_path = self.file_list[index]
        img = Image.open(img_path)

        img_trans = self.transform(
            img, self.phase)

        label = img_path.split('/')[-1].split('.')[0]

        if label == 'cat':
            label = 0
        elif label == 'dog':
            label = 1

        return img_trans, label

#ネコとイヌの画像を分けてロード

train_list_cat = glob.glob('/content/drive/MyDrive/dogs-vs-cats/train/cat.*')
train_list_dog = glob.glob('/content/drive/MyDrive/dogs-vs-cats/train/dog.*')

#学習に使うデータの数を決める

Ntrain = 1000
train_list = train_list_cat[:Ntrain]+train_list_dog[:Ntrain]

#学習データを訓練データと検証データに分けるためのパスのリストを定義

train_idx, valid_idx = train_test_split(range(len(train_list)), test_size = 0.25, random_state = 0)

train_path_list = []
val_path_list = []
for index in train_idx:
    i = train_list[index]
    train_path_list.append(i)

for index in valid_idx:
    i = train_list[index]
    val_path_list.append(i)

# 学習データからテストデータを生成するためのパスのリストを定義
# (Kaggleのテストデータには正解ラベルがないため)

Ntest = 6
test_path_list = train_list_cat[Ntrain:Ntrain+Ntest]+train_list_dog[Ntrain:Ntrain+Ntest]

#　訓練、検証、テストデータの数を確認

len(train_path_list), len(val_path_list), len(test_path_list)

#パスのリストからデータセットを読み込む

size = 150
mean = (0.5,0.5,0.5)
std = (0.5,0.5,0.5)

train_dataset = DogCatDataset(
    file_list = train_path_list, transform = ImageTransform(size, mean, std), phase = 'train')

val_dataset = DogCatDataset(
    file_list = val_path_list, transform = ImageTransform(size, mean, std), phase = 'val')

test_dataset = DogCatDataset(
    file_list = test_path_list, transform = ImageTransform(size, mean, std), phase = 'test')

# 画像が正しく読み込めているか確認

plt.imshow(test_dataset[3][0].permute(1,2,0));

# データをまとめて確認する
plt.figure(figsize=(10, 10))
for n in range(12):
    x, t = test_dataset[n]
    plt.subplot(3, 4, n+1)
    plt.title(test_dataset[n][1])
    plt.axis('off')
    plt.imshow(x.permute(1, 2, 0))

#データローダーでデータをミニバッチに分ける

batch_size = 20
train_dataloader = data.DataLoader(
    train_dataset, batch_size=batch_size, shuffle = True)

val_dataloader = data.DataLoader(
    val_dataset, batch_size=batch_size, shuffle = False)

test_dataloader = data.DataLoader(
    test_dataset, batch_size=batch_size, shuffle = False)

dataloaders_dict = {'train': train_dataloader,
                    'val': val_dataloader}

# モデルの定義
# 4層の隠れ層をもつニューラルネットワークを定義

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(in_channels=3, out_channels=32, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, padding=1)
        self.conv3 = nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, padding=1)
        self.conv4 = nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm2d(32)
        self.bn2 = nn.BatchNorm2d(64)
        self.bn3 = nn.BatchNorm2d(128)
        self.bn4 = nn.BatchNorm2d(256)
        self.fc = nn.Linear(256*9*9, 2)

    def forward(self, x):
        h = self.conv(x)
        h = F.relu(h)
        h = self.bn1(h)
        h = F.max_pool2d(h, kernel_size=2, stride=2)
        h = self.conv2(h)
        h = F.relu(h)
        h = self.bn2(h)
        h = F.max_pool2d(h, kernel_size=2, stride=2)
        h = self.conv3(h)
        h = F.relu(h)
        h = self.bn3(h)
        h = F.max_pool2d(h, kernel_size=2, stride=2)
        h = self.conv4(h)
        h = F.relu(h)
        h = self.bn4(h)
        h = F.max_pool2d(h, kernel_size=2, stride=2)
        h = h.view(-1, 256*9*9)
        h = self.fc(h)
        return h

# モデルのインスタンス化

model = Net()
EPOCHS = 200
loss_fn = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr = 0.0001, momentum=0.9)
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=EPOCHS)

#訓練と検証を行う関数を定義

def train_model(model, dataloaders_dict, loss_fn, optimizer, epochs, scheduler):
    history = {'loss':[], 'acc':[], 'val_loss':[], 'val_acc':[]}
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    model = model.to(device)
    for epoch in range(epochs):
        print(f'{epoch + 1} start')
        print('----------')

        for phase in ['train', 'val']:
            if phase == 'train':
                model.train()
            else:
                model.eval()

            epoch_loss = 0.0
            epoch_corrects = 0.0

            #if (epoch == 0) and (phase == 'train'):
            #    continue

            for inputs, labels in tqdm(dataloaders_dict[phase]):
                inputs = inputs.to(device)
                labels = labels.to(device)

                optimizer.zero_grad()

                with torch.set_grad_enabled(phase == 'train'):
                    outputs = model(inputs)
                    loss = loss_fn(outputs, labels)
                    pred_value, pred_label = torch.max(outputs, 1)

                    if phase == 'train':
                        loss.backward()
                        optimizer.step()

                    epoch_loss += loss.item() * inputs.size(0)
                    epoch_corrects += torch.sum(pred_label == labels.data)



            epoch_loss /= len(dataloaders_dict[phase].dataset)*1.0
            epoch_acc = epoch_corrects.double() / len(dataloaders_dict[phase].dataset)*1.0

            print(f'{phase} Loss: {epoch_loss :.4f} Acc : {epoch_acc :.4f}')
            if phase == 'train':
                history['loss'].append(epoch_loss)
                history['acc'].append(epoch_acc)
            else:
                history['val_loss'].append(epoch_loss)
                history['val_acc'].append(epoch_acc)

            if phase == 'train':
                scheduler.step()
    return history

#学習開始。学習曲線のためにhistoryに格納
%%time
history = train_model(model, dataloaders_dict, loss_fn, optimizer, EPOCHS, scheduler)

# 学習済みモデルの保存
modeldir = '/content/drive/MyDrive/mystudy/20220726_forblog/'
savedparams = 'dogcat_kaggle2000data_mynn_layer2'
torch.save(model.state_dict(), modeldir+savedparams+'.mt')

# 学習曲線を描くための関数を定義

def plot_graph(values1, values2, rng, label1, label2, figname):
    plt.plot(range(rng), values1, label=label1)
    plt.plot(range(rng), values2, label=label2)
    plt.title(os.path.splitext(os.path.basename(figname))[0])
    plt.legend()
    plt.grid()
    plt.savefig(figname)
    plt.show()

#loss curveを描画

t_losses = history['loss']
t_accus = history['acc']
v_losses = history['val_loss']
v_accus = history['val_acc']

lossfigname = '_loss' + '.jpg'

plot_graph(t_losses, v_losses, EPOCHS, 'loss(train)', 'loss(validate)',modeldir+savedparams+lossfigname)

#learning curveを描画

lcfigname = '_learningcurve' + '.jpg'

t_accus_np = [i.to('cpu').detach().numpy().copy() for i in t_accus]
v_accus_np = [i.to('cpu').detach().numpy().copy() for i in v_accus]

plot_graph(t_accus_np, v_accus_np, EPOCHS, 'acc(train)', 'acc(validate)',modeldir + savedparams + lcfigname)

# ネットワークの準備
model = Net().cpu().eval()

# 重みの読み込み
modeldir = '/content/drive/MyDrive/mystudy/20220726_forblog/'
savedparams = 'dogcat_kaggle2000data_mynn_2'
model.load_state_dict(torch.load(modeldir+savedparams+ '.mt', map_location=torch.device('cpu')))

#modelを評価状態にする
model.eval()

#テストデータセットを読み込み
# x : 画像、 t : ラベル(0=ネコ、1=イヌ)
x, t = test_dataset[0]

# NNによる出力値の算出
y = model(x.unsqueeze(0))
print(y)
# --> tensor([[ 1.6474, -1.0637]], grad_fn=<AddmmBackward0>)

# ソフトマックス関数で確率に変換
y = F.softmax(y)
y
# --> tensor([[0.9377, 0.0623]], grad_fn=<SoftmaxBackward0>)　94% ネコ

# テストデータに対する結果(混同行列)を表示

prob = []
pred = []
Y = []
for i, (x,y) in enumerate(test_dataloader):
    with torch.no_grad():
        output = model(x)
    #prob += [i.to('cpu').detach().numpy().copy()[0] for i in output]
    prob += [round(F.softmax(l)[0].to('cpu').detach().numpy().copy()*100,1) for l in output]
    pred += [int(l.argmax()) for l in output]
    Y += [int(l) for l in y]

print(classification_report(Y, pred))

# テストデータに対する予測結果をまとめて表示
plt.rcParams["font.size"] = 18
plt.figure(figsize=(20, 20))
for n in range(len(test_dataset)):
    x, t = test_dataset[n]
    plt.subplot(3, 4,n+1)
    #plt.title(train_dataset[n][1])
    plt.title(str(prob[n])+'%')
    plt.axis('off')
    plt.imshow(x.permute(1, 2, 0))
plt.savefig('/content/drive/MyDrive/mystudy/20220726_forblog/nyans.jpg')

# 自前の画像をロード
mycat_list = glob.glob('/content/drive/MyDrive/mystudy/20220726_forblog/mycat/cat.*')

# リサイズなど画像の前処理
mytest_dataset = DogCatDataset(
    file_list = mycat_list, transform = ImageTransform(size, mean, std), phase = 'test')

mytest_dataloader = data.DataLoader(
    mytest_dataset, batch_size=batch_size, shuffle = False)

# 自前のテストデータに対する結果を表示

prob = []
pred = []
Y = []
for i, (x,y) in enumerate(mytest_dataloader):
    with torch.no_grad():
        output = model(x)
    #prob += [i.to('cpu').detach().numpy().copy()[0] for i in output]
    prob += [round(F.softmax(l)[0].to('cpu').detach().numpy().copy()*100,1) for l in output]
    pred += [int(l.argmax()) for l in output]
    Y += [int(l) for l in y]

print(classification_report(Y, pred))

# 自前テストデータに対する予測結果をまとめて表示
plt.rcParams["font.size"] = 18
plt.figure(figsize=(20, 20))
for n in range(len(mytest_dataset)):
    x, t = mytest_dataset[n]
    plt.subplot(1, 4,n+1)
    #plt.title(train_dataset[n][1])
    plt.title(str(prob[n])+'%')
    plt.axis('off')
    plt.imshow(x.permute(1, 2, 0))
plt.savefig('/content/drive/MyDrive/mystudy/20220726_forblog/mycat.jpg')