PyTorchとResNetで犬と猫を分類するための基本的なプログラムを作成する

今回はPyTorchとResNetを使用して犬と猫を分類するための簡単なプログラムの作成をしていきたいと思います。

なお、環境はDockerを使用し、作成しています。

まずはDockerfileの用意

まずはDockerfileを作成していきます。

# Dockerfile
FROM python:3.8

RUN pip install torch torchvision
RUN pip install scikit-learn

WORKDIR /app

COPY . /app

torchのインストールとscikit-learnも入れています。

Dockerコンテナを立ち上げる

上記のDockerfileがあるディレクトリでコンテナを立ち上げます。

docker build -t pytorch_resnet .

実際のpythonファイル

以下がpytorchとResNetを使用し、犬と猫を分類できるようなpythonファイルのコードです。

import os
import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler
import torchvision
from torchvision import datasets, models, transforms

# Data augmentation and normalization for training
# Just normalization for validation
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),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
}

# adjust the path to your data accordingly
data_dir = '/app/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

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# Load pretrained ResNet
model = models.resnet50(pretrained=True)

# We want to only train the last layer
for param in model.parameters():
    param.requires_grad = False

# Replace the last layer
model.fc = nn.Linear(model.fc.in_features, 2)

model = model.to(device)

criterion = nn.CrossEntropyLoss()

# Observe that all parameters are being optimized
optimizer = optim.SGD(model.fc.parameters(), lr=0.001, momentum=0.9)

# Decay LR by a factor of 0.1 every 7 epochs
exp_lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1)

# print every 10 mini-batches
print_interval = 10

for epoch in range(10):  # loop over the dataset multiple times
    running_loss = 0.0
    for i, data in enumerate(dataloaders['train'], 0):
        # get the inputs; data is a list of [inputs, labels]
        inputs, labels = data[0].to(device), data[1].to(device)

        # zero the parameter gradients
        optimizer.zero_grad()

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

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

print('Finished Training')

# Save the model
torch.save(model.state_dict(), 'model.pth')

では、実際にしていることの解説をしていきます。

データの準備

このコードでは、画像データの前処理を行っています。

画像に対してランダムにリサイズとクロップを行い、ランダムに水平反転を適用し、テンソルに変換し、最後に正規化を行っています。

これは訓練データに対して行います。一方、検証データに対しては、画像をリサイズし、中心をクロップし、テンソルに変換し、最後に正規化を行います。

モデルの定義

次にResNet50モデルをロードし、最終的な全結合層を自分たちのタスク（今回は犬と猫の2分類タスク)に置き換えています。

model = models.resnet50(pretrained=True)

for param in model.parameters():
    param.requires_grad = False

model.fc = nn.Linear(model.fc.in_features, 2)
model = model.to(device)

損失関数と最適化のアルゴリズムの設定

損失関数としてクロスエントロピーロスを使用し、最適化アルゴリズムとしてSGDを使用します。

また、7エポックごとに学習率を0.1倍にするスケジューラも設定しています。

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.fc.parameters(), lr=0.001, momentum=0.9)
exp_lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1)

モデルの訓練

ここでは、モデルを訓練します。訓練ループ内では、まず最適化アルゴリズムの勾配をゼロに設定し、次にモデルを通じてデータを伝播させ、損失を計算し、勾配を計算し、その勾配情報を使って最適化アルゴリズム（ここではSGD）でモデルパラメータを更新します。

そして、そのエポックでの合計損失を追跡します。特定の間隔（ここでは10バッチごと）で、現在のエポックとバッチ番号、そして平均損失を表示します。このプロセスは指定されたエポック数（ここでは10エポック）だけ繰り返されます。

for epoch in range(10):  # loop over the dataset multiple times
    running_loss = 0.0
    for i, data in enumerate(dataloaders['train'], 0):
        # get the inputs; data is a list of [inputs, labels]
        inputs, labels = data[0].to(device), data[1].to(device)

        # zero the parameter gradients
        optimizer.zero_grad()

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

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

モデルの保存

最後に訓練が完了したら、モデルをファイルに保存します。

これによって、後からモデルをロードして予測を行ったり、訓練を再開したりすることが可能となります。

print('Finished Training')

# Save the model
torch.save(model.state_dict(), 'model.pth')

画像ファイルを入れる

続いて、画像ファイルを入れていきます。

犬と猫の画像を以下のようなディレクトリ構成でいれていきます。

../data/
    train/
        dogs/
            dog001.jpg
            dog002.jpg
            ...
        cats/
            cat001.jpg
            cat002.jpg
            ...
    val/
        dogs/
            dog001.jpg
            dog002.jpg
            ...
        cats/
            cat001.jpg
            cat002.jpg
            ...

学習をさせる

では、作ったコードを動かしていきましょう！

docker run -it \
  --name my_container \
  -v /Users/satokitomita/Documents/SFC/resNet/data:/app/data \
  -v /Users/satokitomita/Documents/SFC/resNet:/app \
  pytorch_resnet \
  python main.py

こうすることで無事学習がスタートするはずです。