人工知能に関する断創録

このブログでは人工知能のさまざまな分野について調査したことをまとめています(更新停止: 2019年12月31日)

PyTorch (9) Transfer Learning (Dogs vs Cats)

前回(2018/2/17)は、アリとハチだったけど、今回はイヌとネコでやってみよう。

180209-dogs-vs-cats.ipynb - Google ドライブ

f:id:aidiary:20180221225005j:plain:medium vs. f:id:aidiary:20180221225052j:plain:medium (*^_^*)

import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.optim as optim
from torch.autograd import Variable
import torchvision
from torchvision import datasets, models, transforms

データ準備

Kaggleのイヌとネコのコンペサイトから画像 train.ziptest.zip をダウンロードする。kaggle-cliというコマンドラインツールを使うと、Kaggleにログインした上でコマンドでデータのダウンロードができるのでとても便利!

各種ディレクトリのパスを定義する。

import os
current_dir = os.getcwd()
data_dir = os.path.join(current_dir, 'data', 'dogscats')
train_dir = os.path.join(data_dir, 'train')
valid_dir = os.path.join(data_dir, 'valid')
test_dir = os.path.join(data_dir, 'test')

下のような感じでデータを解凍。fast.ai の受け売りだけどこういうデータ操作もJupyter Notebook上で記録しておくのはよい習慣*1。Jupyter Notebook上でシェルコマンドを使うには ! をコマンドの前につける。

!mkdir $data_dir
!unzip train.zip -d $data_dir
!unzip test.zip -d $data_dir

訓練画像は25000枚、テスト画像は15000枚。訓練データからランダムに選んだ2000枚をバリデーションデータとする。

!mkdir $valid_dir
%cd $train_dir

import os
from glob import glob
import numpy as np
g = glob('*.jpg')
shuf = np.random.permutation(g)
for i in range(2000):
    os.rename(shuf[i], os.path.join(valid_dir, shuf[i]))

PyTorchで読み込みやすいようにクラスごとにサブディレクトリを作成する。Kaggleのテストデータは正解ラベルがついていないため unknown というサブディレクトリにいれる.

# train
%cd $train_dir
%mkdir cats dogs
%mv cat.*.jpg cats/
%mv dog.*.jpg dogs/

# valid
%cd $valid_dir
%mkdir cats dogs
%mv cat.*.jpg cats/
%mv dog.*.jpg dogs/

# test
%cd $test_dir
%mkdir unknown
%mv *.jpg unknown

最終的にデータはこんな階層構造になる。

data
  dogscats
    train
      cats
        cat.1.jpg
        cat.2.jpg
      dogs
        dog.1.jpg
        dog.2.jpg
    valid
      cats
        cat.10.jpg
        cat.21.jpg
      dogs
        dog.8.jpg
        dog.52.jpg
    test
      unknown
        1.jpg
        2.jpg

VGG16のFine-tuning

前回は、ResNet50を使ったので今回はVGG16でやってみた。

vgg16 = models.vgg16(pretrained=True)
print(vgg16)
VGG(
  (features): Sequential(
    (0): Conv2d (3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): ReLU(inplace)
    (2): Conv2d (64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (3): ReLU(inplace)
    (4): MaxPool2d(kernel_size=(2, 2), stride=(2, 2), dilation=(1, 1))
    (5): Conv2d (64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (6): ReLU(inplace)
    (7): Conv2d (128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (8): ReLU(inplace)
    (9): MaxPool2d(kernel_size=(2, 2), stride=(2, 2), dilation=(1, 1))
    (10): Conv2d (128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (11): ReLU(inplace)
    (12): Conv2d (256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (13): ReLU(inplace)
    (14): Conv2d (256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (15): ReLU(inplace)
    (16): MaxPool2d(kernel_size=(2, 2), stride=(2, 2), dilation=(1, 1))
    (17): Conv2d (256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (18): ReLU(inplace)
    (19): Conv2d (512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (20): ReLU(inplace)
    (21): Conv2d (512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (22): ReLU(inplace)
    (23): MaxPool2d(kernel_size=(2, 2), stride=(2, 2), dilation=(1, 1))
    (24): Conv2d (512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (25): ReLU(inplace)
    (26): Conv2d (512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (27): ReLU(inplace)
    (28): Conv2d (512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (29): ReLU(inplace)
    (30): MaxPool2d(kernel_size=(2, 2), stride=(2, 2), dilation=(1, 1))
  )
  (classifier): Sequential(
    (0): Linear(in_features=25088, out_features=4096)
    (1): ReLU(inplace)
    (2): Dropout(p=0.5)
    (3): Linear(in_features=4096, out_features=4096)
    (4): ReLU(inplace)
    (5): Dropout(p=0.5)
    (6): Linear(in_features=4096, out_features=1000)
  )
)

(classifier)(6) の出力が1000クラスになっているのでここをイヌとネコの2クラスにしたい。最初、下のように (6) Linear だけ置き換えられないかトライしてみたがこれはダメなようだ。

num_features = vgg16.classifier[6].in_features
vgg16.classifier[6] = nn.Linear(num_features, 2)  # <= この代入はできない!

1つ上の (classifier) を丸ごと置き換える必要があるみたい。今回は (classifier) より上のレイヤの重みはすべて固定した。(classifier) のみ学習対象とする。

# 全層のパラメータを固定
for param in vgg16.parameters():
    param.requires_grad = False

vgg16.classifier = nn.Sequential(
    nn.Linear(25088, 4096),
    nn.ReLU(inplace=True),
    nn.Dropout(0.5),
    nn.Linear(4096, 4096),
    nn.ReLU(inplace=True),
    nn.Dropout(0.5),
    nn.Linear(4096, 2)    
)

use_gpu = torch.cuda.is_available()
if use_gpu:
    vgg16 = vgg16.cuda()

print(vgg16)

データのロード

まずデータ変換関数を定義する。訓練用とテスト用を分けたが同じ内容でOK。train_preprocess の方にデータ拡張を入れるとさらに精度が上がるかも?

train_preprocess = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

test_preprocess = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

各ディレクトリから画像をロードするDataSetを作成。

train_dataset = datasets.ImageFolder(train_dir, train_preprocess)
valid_dataset = datasets.ImageFolder(valid_dir, test_preprocess)
test_dataset = datasets.ImageFolder(test_dir, test_preprocess)

下のコードでクラスラベルを確認できる。

classes = train_dataset.classes
print(train_dataset.classes)
print(valid_dataset.classes)
print(test_dataset.classes)
['cats', 'dogs']
['cats', 'dogs']
['unknown']

サブディレクトリの cats がラベル0で dogs がラベル1になっていることがわかる。おそらくアルファベット順にラベルが割り当てられるのだろう。unknownはラベル0だけどこれはテスト時にしか使わないので問題ない。

データをミニバッチ単位で取得するDataLoaderを定義。バッチサイズは128にし、訓練データだけシャッフルする。

train_loader = torch.utils.data.DataLoader(train_dataset,
                                           batch_size=128,
                                           shuffle=True)
valid_loader = torch.utils.data.DataLoader(valid_dataset,
                                           batch_size=128,
                                           shuffle=False)

1バッチだけデータを描画してみよう。

def imshow(images, title=None):
    images = images.numpy().transpose((1, 2, 0))  # (h, w, c)
    # denormalize
    mean = np.array([0.485, 0.456, 0.406])
    std = np.array([0.229, 0.224, 0.225])
    images = std * images + mean
    images = np.clip(images, 0, 1)
    plt.imshow(images)
    if title is not None:
        plt.title(title)

images, classes = next(iter(train_loader))
print(images.size(), classes.size())
grid = torchvision.utils.make_grid(images[:25], nrow=5)
imshow(grid)

f:id:aidiary:20180221230323p:plain

モデルの訓練

optimizerには、更新対象のパラメータだけ渡す必要があるので注意。requires_grad = False で固定しているパラメータを含むvgg16.parameters()を指定するとエラーになる。

Fine-tuningなので学習率を小さめにしたSGDを使う。

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(vgg16.classifier.parameters(), lr=0.001, momentum=0.9)

訓練とバリデーションの関数を定義。ここはいつもの。

def train(model, criterion, optimizer, train_loader):
    model.train()
    running_loss = 0
    for batch_idx, (images, labels) in enumerate(train_loader):
        if use_gpu:
            images = Variable(images.cuda())
            labels = Variable(labels.cuda())
        else:
            images = Variable(images)
            labels = Variable(labels)
        
        optimizer.zero_grad()
        outputs = model(images)
        
        loss = criterion(outputs, labels)
        running_loss += loss.data[0]
        
        loss.backward()
        optimizer.step()
    
    train_loss = running_loss / len(train_loader)
    
    return train_loss


def valid(model, criterion, valid_loader):
    model.eval()
    running_loss = 0
    correct = 0
    total = 0
    for batch_idx, (images, labels) in enumerate(valid_loader):
        if use_gpu:
            images = Variable(images.cuda(), volatile=True)
            labels = Variable(labels.cuda(), volatile=True)
        else:
            images = Variable(images, volatile=True)
            labels = Variable(labels, volatile=True)

        outputs = model(images)

        loss = criterion(outputs, labels)
        running_loss += loss.data[0]

        _, predicted = torch.max(outputs.data, 1)
        correct += (predicted == labels.data).sum()
        total += labels.size(0)

    val_loss = running_loss / len(valid_loader)
    val_acc = correct / total
    
    return val_loss, val_acc

出力の logs を予め作成しておく。エポックガンガン回したあとに出力先ディレクトリがなくてエラーになると悲惨・・・あとでTensorBoardXを導入してTensorBoardでロギングするとこういうディレクトリ作成がいらなくなる。

%mkdir logs

Fine-tuningなので5エポックと少しだけ回した。

num_epochs = 5
log_dir = './logs'

best_acc = 0
loss_list = []
val_loss_list = []
val_acc_list = []
for epoch in range(num_epochs):
    loss = train(vgg16, criterion, optimizer, train_loader)
    val_loss, val_acc = valid(vgg16, criterion, valid_loader)

    print('epoch %d, loss: %.4f val_loss: %.4f val_acc: %.4f'
          % (epoch, loss, val_loss, val_acc))
    
    if val_acc > best_acc:
        print('val_acc improved from %.5f to %.5f!' % (best_acc, val_acc))
        best_acc = val_acc
        model_file = 'epoch%03d-%.3f-%.3f.pth' % (epoch, val_loss, val_acc)
        torch.save(vgg16.state_dict(), os.path.join(log_dir, model_file))

    # logging
    loss_list.append(loss)
    val_loss_list.append(val_loss)
    val_acc_list.append(val_acc)
epoch 0, loss: 0.1068 val_loss: 0.0346 val_acc: 0.9905
val_acc improved from 0.00000 to 0.99050!
epoch 1, loss: 0.0351 val_loss: 0.0316 val_acc: 0.9925
val_acc improved from 0.99050 to 0.99250!
epoch 2, loss: 0.0267 val_loss: 0.0315 val_acc: 0.9930
val_acc improved from 0.99250 to 0.99300!
epoch 3, loss: 0.0224 val_loss: 0.0317 val_acc: 0.9930
epoch 4, loss: 0.0175 val_loss: 0.0311 val_acc: 0.9920

バリデーションデータに対する認識率は 99.3%。かなり高い!

テストデータに対する評価

最後にテストデータで評価してみよう。Kaggleのテストデータはラベルがついていないので認識率を求められない。なのでいくつかの画像を描画して目視で確認してみた。

まずは学習済みモデルをロード。PyTorchは重みファイルだけ保存するのが推奨になっていたので、対応するモデル構造は予め用意する必要がある。モデルと重みを両方保存することもできるのかな?

weight_file = 'logs/dogsvscats_ft/epoch003-0.032-0.993.pth'

model = models.vgg16(pretrained=False)
model.classifier = nn.Sequential(
    nn.Linear(25088, 4096),
    nn.ReLU(inplace=True),
    nn.Dropout(0.5),
    nn.Linear(4096, 4096),
    nn.ReLU(inplace=True),
    nn.Dropout(0.5),
    nn.Linear(4096, 2)    
)

model.load_state_dict(torch.load(weight_file,
                                 map_location=lambda storage,
                                 loc: storage))

テストデータのDataLoaderを作成して最初の128枚の画像をロードする。

test_loader = torch.utils.data.DataLoader(test_dataset,
                                          batch_size=128,
                                          shuffle=False)

images, _ = iter(test_loader).next()
images = Variable(images, volatile=True)
print(images.size())

正しく読めてるか最初の25枚を描画して確認しよう。

imshow(torchvision.utils.make_grid(images.data[:25], nrow=5))

OK。test_loadershuffle=False なので順番は固定のはず。でも順番どおりでないな?と思ったら 1.jpg, 10.jpg, 100.jpg, 1000.jpg のようにアルファベット順に並んでいた・・・

f:id:aidiary:20180221231251p:plain

ネットワークに通して予測しよう。

outputs = model(images)
_, predicted = torch.max(outputs.data, 1)

predicted には0(ネコ)、1(イヌ)の予測ラベルが入っている。

print(predicted.numpy())
array([0, 1, 0, 0, 0, 1, 0, 1, 1, 0, 1, 1, 1, 0, 1, 1, 1, 0, 1, 0, 1, 0,
       1, 1, 1, 0, 0, 1, 1, 1, 1, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 1, 1,
       1, 0, 0, 0, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 0, 1, 1, 0, 0, 1])

numpy() してTensorからndarrayに戻したのは出力が上のようにコンパクトに表示されるからで深い意味はない。

この結果から1(イヌ)と予測された画像だけ取り出すには下のように書けばよい。

pred_dogs_images = images[predicted.nonzero().view(-1), :, :, :]
imshow(torchvision.utils.make_grid(pred_dogs_images.data[:25], nrow=5))
  • predicted.nonzero() で0でない(つまり1)のインデックスが取得できる
  • ただし、サイズが (62, 1) のように2D tensorで返ってくるので、view(-1) で1Dtensorにしてからインデキシング

f:id:aidiary:20180221231957p:plain

逆に0(ネコ)と予測された画像だけ取り出すには下のように書けばよい。

pred_cats_images = images[(predicted != 1).nonzero().view(-1), :, :, :]
imshow(torchvision.utils.make_grid(pred_cats_images.data[:25], nrow=5))

f:id:aidiary:20180221232010p:plain

テストデータでもかなりの精度で予測できていることがわかる。

参考

*1:ノートブック再度開いたときにうっかりセル実行しちゃうと大惨事だけど