Docker環境でPyTorch 〜画像解析〜 #02 モデル訓練&保存編

# ディレクトリ構成
~
└── git  
    ├── deep-learning
    │   └── pytorch (★current workspace)
    │       ├── src
    │       │   ├──　car.jpg　（テスト用の画像ファイル）
    │       │   ├──　network.py　（今回作るやつ）
    │       │   ├──　learning.py　（今回作るやつ）
    │       │   └── test.py　（今回作るやつ）
    │       └── docker-compose.yml  
    └── education  
        └── pytorch-examples  
            └── ...

# メタデータ
  batches.meta > クラス名とラベルID(0~9)のマッピング

# トレーニング画像
  data_batch_1 > 10000画像
  data_batch_2 > 10000画像
  data_batch_3 > 10000画像
  data_batch_4 > 10000画像
  data_batch_5 > 10000画像

# テスト画像
  test_batch   > 10000画像

airplane                                        
automobile                                      
bird                                        
cat                                     
deer                                        
dog                                     
frog                                        
horse                                       
ship                                        
truck

import torch.nn as nn
import torch.nn.functional as F

"""
畳み込みニューラルネットワークを定義する
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
ニューラルネットワークセクションからニューラルネットワークをコピーしてから変更します
3チャンネル画像を（1チャンネル画像ではない）定義します。

nn.Conv2d(in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True)
  SEE: https://pytorch.org/docs/stable/nn.html?highlight=linear#torch.nn.Conv2d
  INFO: kernel_sizeで指定した範囲の画像を見て、畳み込みを行っています。画像における畳み込み計算は、簡単にいうとフィルター処理です

nn.MaxPool2d(kernel_size, stride=None, padding=0, dilation=1, return_indices=False, ceil_mode=False)
  SEE: https://pytorch.org/docs/stable/nn.html?highlight=linear#torch.nn.MaxPool2d
  INFO: pool_sizeの範囲を見て、その中で最も大きな値を次の層に渡します。これにより、手書きの線が多少ずれていても、同じような特徴を取り出すことができます。

nn.Linear(in_features, out_features, bias=True)
  SEE: https://pytorch.org/docs/stable/nn.html?highlight=linear#torch.nn.Linear
  INFO: 線形分類器

nn.functional.relu(input, inplace=False)
  SEE: https://pytorch.org/docs/stable/nn.html#torch.nn.functional.relu
  INFO: 整流された線形単位関数を要素ごとに適用します。
"""
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

MODEL_PATH = "model.pth.tar"

import torch
import torchvision
import torchvision.transforms as transforms
import numpy as np
import network
import torch.nn as nn
import torch.optim as optim
import os
from network import MODEL_PATH

"""
CIFAR 10のロードおよび正規化する
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
`torchvision`は、以下のような共通のデータセット用のデータローダーを持っています。
`torchvision.datasets`と `torch.utils.data.DataLoader`は
Imagenet、CIFAR10、MNISTなど、画像用のデータ変換器です。
CIFAR10データセットを使用します。
「飛行機」「自動車」「鳥」「猫」「鹿」「犬」「カエル」「馬」「船」「トラック」などのクラスがあります。
CIFAR-10の画像は、サイズが3×32×32であり、すなわち32×32画素の3チャンネルカラー画像である。

`torchvision.datasets`の出力は範囲[0、1]のPILImage画像です。
正規化された範囲のテンソル[-1,1]に変換します。
"""
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=4, shuffle=True, num_workers=2)

net = network.Net()

"""
損失関数とオプティマイザを定義する
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
分類クロスエントロピー損失とSGDを使用して最適化しましょう。

torch.optim.SGD(params, lr=<object object>, momentum=0, dampening=0, weight_decay=0, nesterov=False)
  SEE: https://pytorch.org/docs/master/optim.html#torch.optim.SGD
  INFO: 最適化アルゴリズム

torch.nn.Module.state_dict(destination=None, prefix='', keep_vars=False)
  SEE: https://pytorch.org/docs/stable/nn.html?highlight=state_dict#torch.nn.Module.state_dict
  INFO: モジュール全体の状態を含む辞書を返します。torch.saveに渡す事で学習パラメータのみの保存します。

torch.optim.Optimizer.state_dict()
  SEE: https://pytorch.org/docs/master/optim.html?highlight=state_dict#torch.optim.Optimizer.state_dict
  INFO: オプティマイザの状態を含む辞書を返します。torch.saveに渡す事で学習パラメータのみの保存します。
"""

# loss function
criterion = nn.CrossEntropyLoss()

# optimizer setting
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

def save_checkpoint(state, filename=MODEL_PATH):
    torch.save(state, filename)

"""
学習済みモデルを使ってさらにチューニング
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
前回学習したモデルがあれば、それを使いさらに学習させましょう。
"""
if os.path.exists(MODEL_PATH):
    source = torch.load(MODEL_PATH, map_location=lambda storage, loc: storage)
    net.load_state_dict(source)

"""
ネットワークを訓練する
^^^^^^^^^^^^^^^^^^^^
ここから面白くなります。データイテレータをループし、入力をネットワークに送り、最適化するだけです。
"""
for epoch in range(2):  # loop over the dataset multiple times
    running_loss = 0.0
    net.train()

    for i, data in enumerate(trainloader, 0):
        # get the inputs
        inputs, labels = data

        # zero the parameter gradients
        optimizer.zero_grad()

        # forward + backward + optimize
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        """
        # 途中経過を毎回保存する事もできます
        save_checkpoint({
            'epoch': epoch + 1,
            'state_dict': net.state_dict(),
            'optimizer' : optimizer.state_dict()
        })
        """

        # print statistics
        running_loss += loss.item()
        if i % 2000 == 1999:    # print every 2000 mini-batches
            print('[%d, %5d] loss: %.3f' %
                  (epoch + 1, i + 1, running_loss / 2000))
            running_loss = 0.0

torch.save(net.state_dict(), MODEL_PATH)
print('Finished Training')

import torch
import torchvision.transforms as transforms
from torch.autograd import Variable #自動微分用
from PIL import Image
import numpy as np
import network
from network import MODEL_PATH

transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

net = network.Net()
net.cpu()
source = torch.load(MODEL_PATH, map_location=lambda storage, loc: storage)
net.load_state_dict(source)

"""
画像をモデルで解析可能な状態のテンソルに変換してネットワークに流し込みます
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
torch.unsqueeze
  SEE: https://pytorch.org/docs/stable/torch.html#torch.unsqueeze
"""
img = Image.open('car.jpg')
img = img.resize((32, 32))
img = transform(img).float()
img = Variable(img)
img = img.unsqueeze(0)

with torch.no_grad():
    outputs = net(img)

    #print(outputs)
    _, predicted = torch.max(outputs, 1)
    print(classes[predicted])