【Python】Pytorchで自作データセットとDataset/DataLoaderを使った画像分類

今回はPytorchを用いた画像分類をPOCで行ったので、その際のメモになります。TesorflowやKerasは以前使ったことがあったので余裕でしょｗｗｗって思っていたら、Pytorch独特の書き方に結構苦戦しました。

<動作環境>

・MAC

・GoogleColab

・Windowの場合は以下からデータセットをダウンロードしてください。

→

Contents

データセットの用意
自作のデータセットを用意する
モデルの準備

モデルにデータセットを学習させる
- 学習済みモデルの保存
- 学習済みモデルの読み込む&検証データでの分類予測

データセットの用意

ファイルの配置のイメージは以下のような感じです。testing trainingにそれぞれ0~9までのラベル名のディレクトリがあり、その中に画像が入っています。

…

自作のデータセットを用意する

この画像データをOPENCVで読み込んでデータセットにします。

まず画像ファイルのパスとラベルを格納したデータフレームを作成します。

# ライブラリの読み込み

from __future__ import print_function

import torch

import torch.nn as nn

import torch.nn.functional as F

import torch.optim as optim

from torch.optim.lr_scheduler import StepLR

from torchvision import datasets

from torchvision import transforms

from torch.utils.data import Dataset

import os

from PIL import Image

from pathlib import Path

from torchvision import transforms

import cv2

import numpy as np

import pandas as pd

data = []

path_ = '/content/mnist_png/training' 

for num in os.listdir(path_):
    for img_path in os.listdir(f'{path_}/{num}')[:100]:
        data.append([f'{path_}/{num}/{img_path}', img_path, num, num])

df = pd.DataFrame(data, columns=['path', 'filename','category', 'label'])

df.head()

このデータフレームを元にデータセットを作るDatasetクラスを定義します。

from torch.utils.data import Dataset

from PIL import Image

class MyDataset(Dataset):

    def __init__(self, train_df, input_size, phase='train',transform=None):
        super().__init__()
        self.train_df = train_df
        image_paths = train_df["path"].to_list()
        self.input_size = input_size
        self.len = len(image_paths)
        self.transform = transform
        self.phase = phase

    def __len__(self):
        return self.len

    def __getitem__(self, index):
        
        image_path = self.train_df["path"].to_list()[index]
        # 画像の読込
        image = cv2.imread(image_path)
        # リサイズ
        image = cv2.resize(image, dsize=(300, 300))
        image = np.array(image).astype(np.float32).transpose(2, 1, 0) # Dataloader で使うために転置する
        # ラベル (0~9)
        label = self.train_df["label"].apply(lambda x : int(x)).to_list()[index]
        return image, label

クラスを定義したら実際にデータフレームを流し込んでインスタンスを作成します。

BATCH_SIZE = 64 

SIZE = 512

image_dataset = MyDataset(
                        df, 
                        (SIZE, SIZE)
                    )

これで自作のデータセットが完成したので、次は訓練用とテスト用に分割します。

# データを7;3で分割する

train_dataset, valid_dataset = torch.utils.data.random_split( image_dataset, [int(len(image_dataset)*0.7), int(len(image_dataset)*0.3)] )

データを分割したらそれぞれのデータをDataloaderに変換します。

from torch.utils.data import DataLoader

# 学習用Dataloader

train_dataloader = DataLoader(
    train_dataset, 
    batch_size=BATCH_SIZE, 
    shuffle=True,
    num_workers=2, 
    drop_last=True,
    pin_memory=True

)

# 評価用Dataloader

valid_dataloader = DataLoader(
    train_dataset, 
    batch_size=BATCH_SIZE, 
    shuffle=True,
    num_workers=2, 
    drop_last=True,
    pin_memory=True
)


# バッチサイズの指定

batch_size = 64

# DataLoaderを作成

train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

valid_dataloader = DataLoader(
    valid_dataset, batch_size=32, shuffle=False)

# 辞書にまとめる

dataloaders_dict = {
    'train': train_dataloader, 
    'valid': valid_dataloader
}

# 動作確認
# イテレータに変換
batch_iterator = iter(dataloaders_dict['train'])

# 1番目の要素を取り出す

inputs, labels = next(batch_iterator)

print(inputs.size())

print(labels)

以下のような実行結果になっていれば成功です。

<実行結果>

torch.Size([64, 3, 300, 300])

tensor([3, 8, 1, 3, 3, 2, 6, 0, 6, 4, 7, 3, 1, 2, 8, 0, 5, 6, 9, 8, 9, 8, 7, 0, 5, 6, 8, 6, 4, 4, 5, 3, 6, 0, 8, 6, 7, 5, 5, 7, 8, 7, 0, 8, 0, 6, 6, 6, 7, 6, 1, 3, 9, 4, 5, 4, 2, 6, 4, 2, 7, 5, 6, 7])

特にlabelsがtensor()型になっているかは注意してください。ここがNumpyarrayやlistだと、tensor型専用メソッドを先の部分で実行した以下のようなエラーが発生します、

ハマったエラー一覧

・img should be PIL Image. Got <class ‘numpy.ndarray’>

・List object has no attribute ‘to’

・’function’ object is not subscriptable

モデルの準備

class Net(nn.Module):
    def __init__(self):
        
        super(Net, self).__init__()
        self.conv1_1 = nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3, padding=1)
        self.conv1_2 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, padding=1)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)

        self.conv2_1 = nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, padding=1)
        self.conv2_2 = nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, padding=1)
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc1 = nn.Linear(in_features=128 * 75 * 75, out_features=128)
        self.fc2 = nn.Linear(in_features=128, out_features=10)

    def forward(self, x):
        x = F.relu(self.conv1_1(x))
        x = F.relu(self.conv1_2(x))
        x = self.pool1(x)

        x = F.relu(self.conv2_1(x))
        x = F.relu(self.conv2_2(x))
        x = self.pool2(x)

        x = x.view(-1, 128 * 75 * 75)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.fc2(x)
        x = F.softmax(x, dim=1)
        return x

net = Net()

print(net)

criterion = nn.CrossEntropyLoss()

optimizer = optim.SGD(net.parameters(), lr=0.01)

nll_loss = nn.NLLLoss()

もしこのコードを他の画像データに置き換えようとして「Target ○ is out of bounds. pytorch～」というようなエラーがでた場合は分類数が恐らくあっていないので上記コードの「out_features=10」の部分を分類数に合わせて変更してください。(今回は0~9の10種類に分類するモデルなので「out_features=10」としています)

モデルにデータセットを学習させる

# エポック数

num_epochs = 10

for epoch in range(num_epochs):

    print('Epoch {}/{}'.format(epoch+1, num_epochs))
    print('-------------')

    for phase in ['train', 'valid']:

        if phase == 'train':
            # モデルを訓練モードに設定
            net.train()
        else:
            # モデルを推論モードに設定
            net.eval()

        # 損失和

        epoch_loss = 0.0

        # 正解数

        epoch_corrects = 0

        # DataLoaderからデータをバッチごとに取り出す
        for inputs, labels in dataloaders_dict[phase]:
            
 
            # optimizerの初期化
            optimizer.zero_grad()

            # 学習時のみ勾配を計算させる設定にする
            with torch.set_grad_enabled(phase == 'train'):

                 outputs = net(inputs)

                 # 損失を計算
                 loss = criterion(outputs, labels)
                
                 # ラベルを予測
                 _, preds = torch.max(outputs, 1)
                
                 # 訓練時はバックプロパゲーション
                 if phase == 'train':

                     # 逆伝搬の計算
                     loss.backward()
                     # パラメータの更新
                     optimizer.step()
                
                # イテレーション結果の計算
                # lossの合計を更新
                # PyTorchの仕様上各バッチ内での平均のlossが計算される。
                # データ数を掛けることで平均から合計に変換をしている。
                # 損失和は「全データの損失/データ数」で計算されるため、
                

                # 平均のままだと損失和を求めることができないため。
                epoch_loss += loss.item() * inputs.size(0)
                
                # 正解数の合計を更新
                epoch_corrects += torch.sum(preds == labels.data)

        # epochごとのlossと正解率を表示

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

        print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc))

学習済みモデルの保存

以下のコードで学習したモデルを保存できます。

# モデルの重みづけを保存する
torch.save(net.state_dict(), "export_model.pth")

学習済みモデルの読み込む&検証データでの分類予測

実際に学習したモデルを読み込んで、検証用の画像を入れて予測してみましょう。

# モデルのインスタンス作成
test_model = Net()

# モデルの重みを読み込む
test_model.load_state_dict(torch.load("export_model.pth"))

# モデルを検証モードに切り替え
test_model.eval()

# 検証用画像のパス
test_image_path = "/content/test.png"

# 訓練用画像と同じように行列処理を行う
img = cv2.imread(test_image_path)

img = cv2.resize(img, dsize=(300, 300))

img = np.array(img).astype(np.float32).transpose(2, 1, 0)

# 一番スコアの高いインデックス番号が返ってくる
results.argmax()

#例 6

Pytorchは自作の画像データデータデータセットを作ろうすると単なるPandasやNumpyのデータ処理ではなくオブジェクト指向的な定義が要求されるので、KerasやTesorflowを使ってきた身からすると少し癖があるように感じました。

慣れたらそうでもなくなりたいと思いたいです。