LoginSignup
41
57

More than 5 years have passed since last update.

Docker環境でPyTorch 〜画像解析〜 #05 最終回 セクシー女優判定編

Last updated at Posted at 2018-10-31

はじめに

株式会社クリエイスCTOの志村とCIOの上田です。
よいよ最終回です。

まずは今までの振り返りをしていきます。

まず最初にdocker環境を構築しpytorchが動くところまでをやってみました。
Docker環境でPyTorch 〜画像解析〜 #01 環境構築編

次にサンプルデータを使ってモデルを訓練し、調教済みモデルを保存して実際に判定までやってみました。
Docker環境でPyTorch 〜画像解析〜 #02 モデル訓練&保存編

次に既に調教されまくったモデルを利用して、少ない学習データからでも精度の高いモデルを保存する転移学習をやりました。
Docker環境でPyTorch 〜画像解析〜 #03 転移学習編

次にセクシー女優を判定すべく学習データを作る処理を行いました。
Docker環境でPyTorch 〜画像解析〜 #04 セクシー女優学習データ作成編

そして最終回の今回は今までやってきたことをまとめて、最終形態のコードを作って実際に判定してみたいと思います。

前提条件

  1. Docker環境でPyTorch 〜画像解析〜 #01〜#04までを終わらせていること
  2. 完全に趣味の世界なので、楽しんで行ってください

流れ

それでは、ざっくりとこんな感じで進めて行きます
1. 学習データを読み込む為の処理と判定するときのラベルリスト作成処理だけを分離したdatasets.pyを作成
2. #3で作ったsample.pyをちょっと改造して転移学習用learning.pyを作成
3. #2で作ったtest.pyをちょっと改造して判定用api.pyを作成
4. learning.pyを実行してモデルを訓練
5. api.pyを実行してセクシー女優判定

# ディレクトリ構成図
~
└── git  
    ├── deep-learning
    │   └── pytorch (★current workspace)
        │        ├── data
    │       │   └──av_data 
       │       │       └── transfer 
    │       │           ├── train (教師画像 face_detect.py実行で配下は自動作成されます)
    │       │           │    ├── JULIA
    │       │           │    │    └── xxx.jpg (JULIA様の顔をトリミングした64×64の画像100枚程度)
    │       │           │    ├── Rio
    │       │           │    ⁝    └── xxx.jpg (Rio様の顔をトリミングした64×64の画像100枚程度)
    │       │           └── val (検証画像 今回は使いませんが以下のフォルダ構成になる予定)
    │       │                ├── JULIA
    │       │                │    └── xxx.jpg (サイズがバラバラのJULIA様の画像70枚程度)
    │       │                ├── Rio
    │       │                ⁝     └── xxx.jpg (サイズがバラバラのRio様の画像70枚程度)   
    │       ├── src
    │       │   └── execute
       │       │              ├── api.py (今回作るやつ)
       │       │              ├── datasets.py (今回作るやつ)
       │       │              └── learning.py (今回作るやつ)
    │       └── docker-compose.yml  
    └── ...

1. 学習データを読み込む為の処理と判定するときのラベルリスト作成処理

datasets.py
import torch
import torchvision.transforms as transforms
from torchvision import datasets
import os

# データ変換関数
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/av_data/transfer'
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']}
classes = image_datasets['train'].classes
label_list_num = len(classes)

2. 転移学習処理

learning.py
import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler
from torchvision import models
from datasets import dataloaders, dataset_sizes, label_list_num

import time
import copy
import numpy as np

# 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':
                    backward + optimize
                    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層の線形レイヤー入力数を女優数に置き換える
num_features = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_features, label_list_num)

# 目的関数とオブティマイザを作成
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')

3. セクシー女優判定処理

api.py
import torch
import torch.nn as nn
from torch.autograd import Variable
from PIL import Image
from torchvision import models
from datasets import data_transforms, classes, label_list_num


model_ft = models.resnet18(pretrained=True)
num_features = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_features, label_list_num)

model_ft.cpu()
source = torch.load('model_ft.pkl', map_location=lambda storage, loc: storage)
model_ft.load_state_dict(source)

# 判定したい画像パスを入れてください
img = Image.open('/workspace/data/av_data/transfer/val/***/***.jpg')
img = data_transforms['val'](img).float()
img = Variable(img)
img = img.unsqueeze(0)

model_ft.eval()
with torch.no_grad():
    outputs = model_ft(img)
    _, predicted = torch.max(outputs.data, 1)
    print(classes[predicted])

4. モデルを転移学習させる

まずはdockerを起動して中に入ります

docker-compose run --name pytorch.examples --entrypoint=/bin/bash --rm pytorch.examples
python /workspace/src/execute/learning.py

5. よいよセクシー女優判定

python /workspace/src/execute/api.py

ちゃんと判定できましたでしょうか?!
うまくいかなかった場合はもう少しEpochを増やして学習させてみるとか、重み付けを固定したりしてみると精度が上がるかもしれません。
5ヶ月間に渡り連載してきましたが、終わってみれば案外あっさりしたコードだと思いませんか?
皆さま良いAIライフを!

ご視聴ありがとうございました。

41
57
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
41
57