PyTorch で簡単なニューラルネットワークを構築し画像を分類する方法

シンプルな CNN

PyTorch は深層学習用のフレームワークの一つである。PyTorch を用いることで、様々なニューラルネットワークを容易に構築できる。このページでは、画像分類を行うための畳み込みニューラルネットワークを PyTorch で作成して、学習と検証を行う例を示す。

画像データセット

サンプル画像として Kaggle にある Blood Cell Images を使用する。このデータセットをダウンロードして展開すると、dataset-master と dataset2-master の二つのフォルダが得られる。ここでは、dataset2-master フォルダの下の images フォルダの画像を使用する。

ls dataset2-master/dataset2-master/images
## TEST        TEST_SIMPLE TRAIN

ls dataset2-master/dataset2-master/images/TRAIN
## EOSINOPHIL LYMPHOCYTE MONOCYTE   NEUTROPHIL

ls dataset2-master/dataset2-master/images/TEST
## EOSINOPHIL LYMPHOCYTE MONOCYTE   NEUTROPHIL

このデータセットは 4 つのカテゴリに分けられている。これから、入力された画像を EOSINOPHIL、LYMPHOCYTE、MONOCYTE、および NEUTROPHIL のいずれかに分類するモデルを作成する。

以上のように、データをダウンロードし、データのフォルダ構造を確認できたのちに、これを PyTorch で、プログラムがわかるような形で記述する。これには以下のような手順を踏む。

  1. 画像フォルダ中の画像を確認し、学習可能な形で Python オブジェクトに保存する(ImageFolder)。
  2. 画像データ前処理
    • 畳み込みニューラルネットワークによる画像分類では、入力画像のサイズを揃える必要がある。ここでは、すべての入力画像を幅 320 ピクセル、縦 240 ピクセルに揃えることにする(transforms.Resize)。
    • PyTorch では、行列データとして読み込まれた画像データをそのまま利用できないので、これをテンソル型に変換する(transforms.ToTensor)。
    • また、画像データは 0 から 255 までの整数によって記述されている。このままで学習を進めるには数値が多すぎて収束しにくいため、これらの数値を小さな数値に変換するための正規化を行う(transforms.Normalize)。
  3. 画像フォルダの画像のなから数まいずつ取り出して、前処理を行い、学習可能な形で用意する(DataLoader)。
data_dir_path = 'dataset2-master/dataset2-master/images/'

data_transforms = torchvision.transforms.Compose([
    torchvision.transforms.Resize((240, 320)),
    torchvision.transforms.ToTensor(),
    torchvision.transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

train_dataset = torchvision.datasets.ImageFolder(root=data_dir_path + 'TRAIN',
                                                 transform=data_transforms)
train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=4,
                                               shuffle=True, num_workers=1)

batch_size=4 と指定することで、DataLoader は写真を 4 枚ずつ読み込んで学習に回すことができるようになる。

ニューラルネットワークの定義

PyTorch の nn.Module クラスを継承してニューラルネットワークの構造を定義する。ここでは、畳み込み層、プーリング層、畳み込み層、プーリング層、全結合層、全結合層、そして出力層のような簡単なアーキテクチャを定義する。

import torch
import torchvision

class Net(torch.nn.Module):

    def __init__(self):
        super(Net, self).__init__()

        self.conv1 = torch.nn.Conv2d(3, 16, 5)
        self.pool1 = torch.nn.MaxPool2d(2, 2)
        self.conv2 = torch.nn.Conv2d(16, 32, 5)
        self.pool2 = torch.nn.MaxPool2d(2, 2)

        self.fc1 = torch.nn.Linear(32 * 57 * 77, 512)
        self.fc2 = torch.nn.Linear(512, 64)
        self.fc3 = torch.nn.Linear(64, 4)


    def forward(self, x):

        x = torch.nn.functional.relu(self.conv1(x))
        x = self.pool1(x)
        x = torch.nn.functional.relu(self.conv2(x))
        x = self.pool2(x)
        #print(x.shape)
        x = x.view(-1, 32 * 57 * 77)
        #print(x.shape)
        x = torch.nn.functional.relu(self.fc1(x))
        x = torch.nn.functional.relu(self.fc2(x))

        x = self.fc3(x)

        return x

畳み込み層およびプーリング層の演算結果を全結合層に引き渡す時、複数枚の行列データを一つのベクトルとして表示させる必要があるので、ここでは NumPy の view メソッドを使用してデータ配置の変換を行う。上の例では、view メソッドの第 2 引数に 32 * 57 * 77 を指定したが、これは入力カラー画像が 3x320x240 で、畳み込み演算とプーリング演算を繰り返した結果 32x47x77 となっているためである。

なお、第 2 引数の数値は、畳み込み演算とプーリング演算を追跡すれば、正確に計算することができる。層が深くなり、計算が煩雑になる場合、view メソッドの直前に print(x.shape) を実行すると、それまでのデータ構造が出力される。

学習

ミニバッチ学習できるように DataLoader の用意ができ、ニューラルネットワークのアーキテクチャの定義もできた。ここからデータをネットワークに代入して学習を進める。学習を進めるとき、以下のような手順で行う。

  1. 損失関数(CrossEntropyLoss)、重みの更新アルゴリズムおよび学習率(SGD(lr=0.01))を定義する。
  2. 指定されたエポック数(n_epochs)だけ以下の計算を繰り返す。
    1. 各ミニバッチに対して以下の計算を繰り返す。
      1. ネットワークに画像を代入して、予測結果を得る。
      2. 予測結果と教師ラベル(フォルダ名)を比較して、損失を計算する。
      3. 損失をネットワークの前方に伝播し、重みを更新する(backward)。
net = Net()

criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(net.parameters(), lr=0.01, momentum=0.9)

n_epochs = 50
for epoch in range(n_epochs):
    # loss in this epoch
    running_loss = 0.0

    for inputs, labels in train_dataloader:        
        optimizer.zero_grad()
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()

    print('Epoch: {};  Loss: {}'.format(epoch, running_loss))

print('Finished Training')

検証・予測

訓練データで学習を済ませたモデルにテストデータを代入して検証する。テストデータも、訓練データと同様にミニバッチでテストできる。

test_dataset = torchvision.datasets.ImageFolder(root=data_dir_path + 'TEST',
                                                transform=data_transforms)
test_dataloader = torch.utils.data.DataLoader(test_dataset, batch_size=4,
                                              num_workers=1)


n_correct = 0
n_total = 0

with torch.no_grad():
    for data in test_dataloader:
        images, labels = data
        outputs = net(images)
        _, predicted = torch.max(outputs.data, 1)
        n_total += labels.size(0)
        n_correct += (predicted == labels).sum().item()

print('#images: {}; acc: {}'.format(n_total, 100 * n_correct / n_total))