Docker環境でPyTorch 〜画像解析〜 #03 転移学習編

はじめに

株式会社クリエイスCTOの志村です。
前回の続きです。
この記事に最初に行き着いた方は前回の記事を見ていただき、環境を作るところから始めてください。
Docker環境でPyTorch 〜画像解析〜 #01 環境構築編

volumesに- ./data:/workspace/dataを追記しました

今回も上田先生 に解説してもらいます。

社員の曽宮、モトキが四苦八苦しながらまとめているこちらの記事もよろしかったらどうぞ
Anaconda環境でPyTorch 〜株価予想〜 #03 予測編

この記事でやること

事前調教済みモデルのResNet18を使って転移学習を行います。

この記事でやらないこと

1. Dockerfileを使って公式Imageを元にカスタムImageを作りません
2. AWSのECS,　EKSには対応いたしません
3. コピペしやすい様にコードブロックの頭のユーザーマーク($, %, #)は記述いたしません
4. ファインチューニングは行いません

前提条件

1. 実験場を作成する目的
2. 今回は最終にImageをビルドする訳ではない
3. というわけで、かな〜り緩めにvolumesで実験ソースをマウントして行きます
4. ホストでIDEをつかってガンガン修正しつつ、コンテナ上でガンガンテストできる環境を目指します

流れ

それでは、ざっくりとこんな感じで進めて行きます

1. ResNet18とは
2. 教師データ(蜂&蟻画像)をダウンロードしディレクトリに配置
3. テストコードを作る
4. 動かして確認する

Localに次の様なディレクトリ構造で構築して行きます
前回の記事からの流れになっているので階層深めです

# ディレクトリ構成図
~
└── git  
    ├── deep-learning
    │   └── pytorch (★current workspace)
　　　　　　　　│　      　├── data
    │       │   └──　hymenoptera_data　
    │       │   　　   ├── train　(教師画像)
    │       │   　　   │    ├── ants
    │       │   　　   │    │    └── xxx.jpg　(サイズがバラバラの蟻120枚)
    │       │   　　   │    └── bees
    │       │   　　   │    　　    └── xxx.jpg　(サイズがバラバラの蜂120枚)
    │       │   　　   └── val　(検証画像)
    │       │   　　   　　    ├── ants
    │       │   　　   　　    │    └── xxx.jpg　(サイズがバラバラの蟻75枚)
    │       │   　　   　　    └── bees
    │       │   　　   　　    　　    └── xxx.jpg　(サイズがバラバラの蜂75枚)
    │       ├── src
    │       │   └── transfer_learning
 　　　　　　│       │　　　　　　　　　　　　　　└── sample.py　(今回作るやつ)
    │       └── docker-compose.yml  
    └── ...

1. ResNet18とは

画像を1000クラス分類する事ができる訓練済みのモデルで
224x224の画像を分類する事ができます。

pytorchでは他にこんなモデルが利用できるみたいですよ

• AlexNet
• VGG
• ResNet
• SqueezeNet
• DenseNet
• Inception v3

詳しい詳細はこちら

2. 教師データ(蜂&蟻画像)をダウンロードしディレクトリに配置

cd ~/git/deep-learning/pytorch
wget https://download.pytorch.org/tutorial/hymenoptera_data.zip -P data/
unzip data/hymenoptera_data.zip -d data

データの中身は前述した「ディレクトリ構成図」のhymenoptera_data配下の通りです。

3. テストコードを作る

sample.py

import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler
from torch.autograd import Variable
import torchvision
from torchvision import datasets, models, transforms
import os
import time
import copy
import numpy as np


# データ変換関数

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, 256)),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
}


# DataSetをディレクトリから画像を直接読み込む
data_dir = '/workspace/data/hymenoptera_data'
image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x), data_transforms[x])
                  for x in ['train', 'val']}
dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x],
                                              batch_size=4,
                                              shuffle=True,
                                              num_workers=4) for x in ['train', 'val']}
dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']}
class_names = image_datasets['train'].classes

# macのdockerはGPUが物理的に利用できないのでCPUで動きます(怒)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# 訓練用関数

def train_model(model, criterion, optimizer, scheduler, num_epochs=25):
    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, num_epochs - 1))
        print('-' * 10)

        # 各エポックで訓練+バリデーションを実行
        for phase in ['train', 'val']:
            if phase == 'train':
                scheduler.step()
                model.train   # training mode
            else:
                model.eval  # evaluate mode

            running_loss = 0.0
            running_corrects = 0

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

                # zero the parameter gradients
                optimizer.zero_grad()

                # forward
                outputs = model(inputs)
                _, preds = torch.max(outputs.data, 1)
                loss = criterion(outputs, labels)

                if phase == 'train':
                    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 == 'val' 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



# resnet18ロード
model_ft = models.resnet18(pretrained=True)

# 出力がImageNetの1000クラス分類なので (Linear(512, 1000)) fc層を2クラス(蟻、蜂)に置き換える
num_features = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_features, 2)

criterion = nn.CrossEntropyLoss()
optimizer_ft = optim.SGD(model_ft.parameters(), lr=0.001, momentum=0.9)

# 7エポックごとに学習率を0.1倍する
exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=7, gamma=0.1)

model_ft = train_model(model_ft, criterion, optimizer_ft, exp_lr_scheduler, num_epochs=25)
# 調教済みモデルをファイル出力して終わり
torch.save(model_ft.state_dict(), 'model_ft.pkl')

4. 動かして確認する

docker rm -f pytorch
docker-compose run --name pytorch --entrypoint=/bin/bash pytorch
cd /workspace/src/transfer_learning
python sample.py

出力はこんな感じ

Epoch 0/24
----------
train Loss: 0.6449 Acc: 0.6885
val Loss: 0.7589 Acc: 0.6928

Epoch 1/24
----------
train Loss: 0.5536 Acc: 0.7541
val Loss: 0.3716 Acc: 0.8235

Epoch 2/24
----------
.
.
.
.
.

なんか出た..... コーヒー飲みながら25エポック待ちます!!!

こんな感じで意外とお手軽に学習済みモデルを利用することができました。
あとはファインチューニングを行えば立派なAIとして動くことでしょう！

次回

キャプチャ認証の解析に向けて教師データを作っていこうかと思ってたのですが、googleが今後違う認証使って来たらダルいのと、ちょっと熱が冷めてしまったというのもあって
####AV女優の顔、体から名前を当てる的なやつ
をやってみたい！

ひとまず、みんなに相談して決めたいと思います。。