LoginSignup
3
3

More than 1 year has passed since last update.

@john-rocky(MLBOYまじま だいすけお仕事大歓迎)inエンジニアと人生コミュニティ

カスタムデータセットでなるべくかんたんに画像分類器をつくりたい。Pytorchで転移学習

Last updated at Posted at 2022-04-14

自前のデータセットでImageClassificationモデルをトレーニングします

Pytorch　チュートリアルのColab:

ユースケースに合わせて、画像分類機能をつくりたい

画像分類を自身のユースケースに合わせて行いたい場合、ユースケースにあったデータセットで画像分類モデルをトレーニングします。
たとえば、画像内の交通標識を見分けたい場合は、それぞれの標識の画像を集めて、モデルに学習してもらいます。

Pytorchを使ってやってみよう

Pytorchを使って学習をしてみます。
今回は事前にImageNetというデータセットを学習したモデルを使って、モデルの最後に未学習のニューラルネットワークをつけ、学習させます。
この手法はFinetuneと呼ばれます。
また、記事の最後では、事前学習済みレイヤーはロックして、後者の付け足した最後の層のみ学習させます。
モデル全体をランダムの重みから学習させるのに対して、この手法は転移学習と呼ばれます。
大規模なデータセットを用意できないユースケースでは一般的な手法です。

方法

１、必要なモジュールのインポート

from __future__ import print_function, division

import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler
import torch.backends.cudnn as cudnn
import numpy as np
import torchvision
from torchvision import datasets, models, transforms
import matplotlib.pyplot as plt
import time
import os
import copy

cudnn.benchmark = True
plt.ion()   # interactive mode

２、データセットを用意します

以下のように、データのディレクトリ内にtrainとvalのディレクトリを用意し、それぞれ中にクラス名のディレクトリを用意して、画像を入れます。

dataset
|
|__train
|  |
|  |__class1
|  |  |
|  |  |__1.jpg
|  |  |__2.jpg
|  |  |__...
|  |
|  |__class2
|  |  |
|  |  |__1.jpg
|  |  |__2.jpg
|  |  |__...
|  |
|  |__...
|
|__val
|  |
|  |__class1
|     |
|     |__1.jpg
|     |__2.jpg
|...

３、データセットをモデルに与えるデータローダーの形式にします

# データセットの前処理を設定します
data_transforms = {
    'train': transforms.Compose([
        transforms.RandomResizedCrop(224), # ランダムにリサイズしてクロップします
        transforms.RandomHorizontalFlip(), # 水平方向にフリップします
        transforms.ToTensor(), # torch用のテンソルにします
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) # ImageNet形式の正規化をします
    ]),
    '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])
    ]),
}

data_dir = 'dataset'

# datasetというモジュールのImageFolderという関数に、ラムダ関数でデータセットを与えて、上述の前処理を加え、torch用のデータセット形式にする
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")

データローダー形式にしたデータを画像表示して、確認します。

def imshow(inp, title=None):
    """Imshow for Tensor."""
    inp = inp.numpy().transpose((1, 2, 0)) # テンソルをnumpy arrayに
    # 非正規化
    mean = np.array([0.485, 0.456, 0.406]) 
    std = np.array([0.229, 0.224, 0.225])
    inp = std * inp + mean
    inp = np.clip(inp, 0, 1)
    # 表示
    plt.imshow(inp)
    if title is not None:
        plt.title(title)
    plt.pause(0.001)  # pause a bit so that plots are updated


# データローダーから読み込む
inputs, classes = next(iter(dataloaders['train']))

# グリッド画像にする
out = torchvision.utils.make_grid(inputs)

imshow(out, title=[class_names[x] for x in classes])

トレーニングループを設定する

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(f'Epoch {epoch}/{num_epochs - 1}')
        print('-' * 10)

        # Each epoch has a training and validation phase
        for phase in ['train', 'val']:
            if phase == 'train':
                model.train()  # Set model to training mode
            else:
                model.eval()   # Set model to evaluate mode

            running_loss = 0.0
            running_corrects = 0

            # Iterate over data.
            for inputs, labels in dataloaders[phase]:
                inputs = inputs.to(device)
                labels = labels.to(device)

                # zero the parameter gradients
                optimizer.zero_grad()

                # forward
                # track history if only in train
                with torch.set_grad_enabled(phase == 'train'):
                    outputs = model(inputs)
                    _, preds = torch.max(outputs, 1)
                    loss = criterion(outputs, labels)

                    # backward + optimize only if in training phase
                    if phase == 'train':
                        loss.backward()
                        optimizer.step()

                # statistics
                running_loss += loss.item() * inputs.size(0)
                running_corrects += torch.sum(preds == labels.data)
            if phase == 'train':
                scheduler.step()

            epoch_loss = running_loss / dataset_sizes[phase]
            epoch_acc = running_corrects.double() / dataset_sizes[phase]

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

            # 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(f'Training complete in {time_elapsed // 60:.0f}m {time_elapsed % 60:.0f}s')
    print(f'Best val Acc: {best_acc:4f}')

    # load best model weights
    model.load_state_dict(best_model_wts)
    return model

モデルとトレーニング設定を定義します

torchvisionの事前トレーニング済みモデルをダウンロードし、トレーニングするニューラルネットワークを付け足します。
ロスファンクションとオプティマイザーを設定します。

model_ft = models.resnet18(pretrained=True) # 事前学習済みモデル
num_ftrs = model_ft.fc.in_features # 事前学習済みモデルのファインチューン用のフィーチャー出力

# 事前トレーニング済みモデルの出力形状の出力を入力として受け取り、今回分類するクラス数分の出力をするニューラルネットワークをファインチューンレイヤーとして付け足す
model_ft.fc = nn.Linear(num_ftrs, len(class_names)) 

model_ft = model_ft.to(device)

criterion = nn.CrossEntropyLoss() # ロス関数

# Observe that all parameters are being optimized
optimizer_ft = optim.SGD(model_ft.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_ft, step_size=7, gamma=0.1)

事前トレーニング済みモデルは、torchvision.modelsでは以下が提供されています。

import torchvision.models as models
resnet18 = models.resnet18(pretrained=True)
alexnet = models.alexnet(pretrained=True)
squeezenet = models.squeezenet1_0(pretrained=True)
vgg16 = models.vgg16(pretrained=True)
densenet = models.densenet161(pretrained=True)
inception = models.inception_v3(pretrained=True)
googlenet = models.googlenet(pretrained=True)
shufflenet = models.shufflenet_v2_x1_0(pretrained=True)
mobilenet_v2 = models.mobilenet_v2(pretrained=True)
mobilenet_v3_large = models.mobilenet_v3_large(pretrained=True)
mobilenet_v3_small = models.mobilenet_v3_small(pretrained=True)
resnext50_32x4d = models.resnext50_32x4d(pretrained=True)
wide_resnet50_2 = models.wide_resnet50_2(pretrained=True)
mnasnet = models.mnasnet1_0(pretrained=True)
efficientnet_b0 = models.efficientnet_b0(pretrained=True)
efficientnet_b1 = models.efficientnet_b1(pretrained=True)
efficientnet_b2 = models.efficientnet_b2(pretrained=True)
efficientnet_b3 = models.efficientnet_b3(pretrained=True)
efficientnet_b4 = models.efficientnet_b4(pretrained=True)
efficientnet_b5 = models.efficientnet_b5(pretrained=True)
efficientnet_b6 = models.efficientnet_b6(pretrained=True)
efficientnet_b7 = models.efficientnet_b7(pretrained=True)
regnet_y_400mf = models.regnet_y_400mf(pretrained=True)
regnet_y_800mf = models.regnet_y_800mf(pretrained=True)
regnet_y_1_6gf = models.regnet_y_1_6gf(pretrained=True)
regnet_y_3_2gf = models.regnet_y_3_2gf(pretrained=True)
regnet_y_8gf = models.regnet_y_8gf(pretrained=True)
regnet_y_16gf = models.regnet_y_16gf(pretrained=True)
regnet_y_32gf = models.regnet_y_32gf(pretrained=True)
regnet_x_400mf = models.regnet_x_400mf(pretrained=True)
regnet_x_800mf = models.regnet_x_800mf(pretrained=True)
regnet_x_1_6gf = models.regnet_x_1_6gf(pretrained=True)
regnet_x_3_2gf = models.regnet_x_3_2gf(pretrained=True)
regnet_x_8gf = models.regnet_x_8gf(pretrained=True)
regnet_x_16gf = models.regnet_x_16gf(pretrainedTrue)
regnet_x_32gf = models.regnet_x_32gf(pretrained=True)
vit_b_16 = models.vit_b_16(pretrained=True)
vit_b_32 = models.vit_b_32(pretrained=True)
vit_l_16 = models.vit_l_16(pretrained=True)
vit_l_32 = models.vit_l_32(pretrained=True)
convnext_tiny = models.convnext_tiny(pretrained=True)
convnext_small = models.convnext_small(pretrained=True)
convnext_base = models.convnext_base(pretrained=True)
convnext_large = models.convnext_large(pretrained=True)

学習

model_ft = train_model(model_ft, criterion, optimizer_ft, exp_lr_scheduler,
                       num_epochs=25)

Best val Acc: 0.934641

２クラス200枚程度、ColabのCPU（GPUなし）で33分で学習できました。

モデルで推論してみる

def visualize_model(model, num_images=6):
    was_training = model.training
    model.eval()
    images_so_far = 0
    fig = plt.figure()

    with torch.no_grad():
        for i, (inputs, labels) in enumerate(dataloaders['val']):
            inputs = inputs.to(device)
            labels = labels.to(device)

            outputs = model(inputs)
            _, preds = torch.max(outputs, 1)

            for j in range(inputs.size()[0]):
                images_so_far += 1
                ax = plt.subplot(num_images//2, 2, images_so_far)
                ax.axis('off')
                ax.set_title(f'predicted: {class_names[preds[j]]}')
                imshow(inputs.cpu().data[j])

                if images_so_far == num_images:
                    model.train(mode=was_training)
                    return
        model.train(mode=was_training)

visualize_model(model_ft)

ダウンロード (46).png
ダウンロード (47).png

最後の層以外はロックする(固定特徴抽出器)

学習がより高速になります。

トレーニングデータが事前トレーニングされたネットワークがトレーニングされたデータと本質的に類似している場合は、固定特徴抽出器の上で新しい分類器をトレーニングするだけで十分かもしれません（そして間違いなく最速のアプローチです）。ただし、データがまったく異なる場合（たとえば、絵画やイラストを微調整しようとしているときにconv netが写真でトレーニングされた場合）、以前のレイヤーも微調整するのが理にかなっている場合があります。
https://forums.fast.ai/t/two-types-of-transfer-learning/9100

model_conv = torchvision.models.resnet18(pretrained=True)
for param in model_conv.parameters():
    param.requires_grad = False #　学習をロック

# Parameters of newly constructed modules have requires_grad=True by default
num_ftrs = model_conv.fc.in_features
model_conv.fc = nn.Linear(num_ftrs, 2)

model_conv = model_conv.to(device)

criterion = nn.CrossEntropyLoss()

# Observe that only parameters of final layer are being optimized as
# opposed to before.
optimizer_conv = optim.SGD(model_conv.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_conv, step_size=7, gamma=0.1)

model_conv = train_model(model_conv, criterion, optimizer_conv,
                         exp_lr_scheduler, num_epochs=25)

Best val Acc: 0.967320

２クラス200枚程度、ColabのCPU（GPUなし）で16分で学習できました。

ユースケースに合わせた判別器を作ろう

データを揃えれば、短時間で学習できるので、ビジネスなどに応用しましょう。

🐣

フリーランスエンジニアです。
お仕事のご相談こちらまで
rockyshikoku@gmail.com

Core MLやARKitを使ったアプリを作っています。
機械学習／AR関連の情報を発信しています。

Twitter
Medium
GitHub

3
3
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
Sign upLogin
3
3