PyTorch: 事前学習モデルを活用した画像分類のファインチューニング

1. はじめに

本ブログでは、PyTorchを使用して事前学習済みモデルを活用し、画像分類タスクでそのパフォーマンスを向上させる方法について説明します。具体的な実装手順も紹介します。

2. 事前学習モデルとは

事前学習モデル(pre-trained model)とは、大規模なデータセットで訓練されたニューラルネットワークのことです。これらのモデルは特定のタスクに対する知識を持っており、新しいタスクでそれらを再利用することが可能です。

3.PyTorchとは

PyTorchはオープンソース機械学習ライブラリであり、Pythonプログラム内から使用することが可能です。テンソル計算（NumPyに似ていますがGPU上でも動作します）や深層学習アプリケーション開発・運用等に役立ちます。

4 ファインチューニング（Fine-tuning）とは

ファインチューニング（Fine-tuning）では、あるタスクで訓練された既存のモデルを、新しいタスクに対応するように調整します。これは転移学習の一種であり、ベースとなるモデルが新しいタスクに有用な特徴を抽出する能力を既に持っているという考え方に基づいています。

5. 実装手順

5.1 必要なライブラリのインポート

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

5.2 データセットの準備

ここではtorchvision.datasetsからCIFAR10データセットを使用します。

transform = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

trainset = datasets.CIFAR10(root='./data', train=True,
                                        download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=4,
                                          shuffle=True)

testset = datasets.CIFAR10(root='./data', train=False,
                                       download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=4,
                                         shuffle=False)

5.3事前学習モデルのロードと調整

ResNet18という事前学習モデルを使用して、新しいタスクに適応させるための調整を行います。

model = models.resnet18(pretrained=True)

# モデルの最終層の出力ユニット数を新しいタスクに合わせて変更
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, num_classes)  # num_classesは新しいタスクのクラス数

# GPUが利用可能な場合はモデルをGPU上に配置する
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = model.to(device)

5.4 トレーニングと評価

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

# モデルのトレーニングと評価（エポック数やバッチサイズは適宜変更してください）
num_epochs = 10

for epoch in range(num_epochs):
    running_loss = 0.0
    
    for i, data in enumerate(trainloader, 0):
        inputs, labels = data[0].to(device), data[1].to(device)
        
        optimizer.zero_grad()
        
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        
        running_loss += loss.item()
        
    print(f"Epoch {epoch+1} - Loss: {running_loss / len(trainloader)}")
    
# テストデータでの精度評価
correct_count = 0
total_count = 0

with torch.no_grad():
    for data in testloader:
        images, labels = data[0].to(device), data[1].to(device)
        
        outputs = model(images)
        _, predicted_labels = torch.max(outputs.data, 1)
        
        total_count += labels.size(0)
        correct_count += (predicted_labels == labels).sum().item()

accuracy = correct_count / total_count * 100
print(f"Test Accuracy: {accuracy}%")

結論及び今後の方針

本ブログでは、PyTorchを使用して事前学習済みモデルを活用した画像分類タスクのファインチューニング手法について紹介しました。事前学習モデルを利用することで、少ないデータセットでも高い性能が期待できます。今後はさらに様々な画像分類タスクにおけるファインチューニング手法やパフォーマンス向上策についても探求していく予定です。

AIko Code Symphony