【PyTorch入門】オリジナル画像分類CNNモデルの構築と学習パイプラインの実装

プログラミング

2026.03.31

OpenCVのアプローチで前処理を学んだあとは、自らの手でAI（ニューラルネットワーク）を設計し、独自の画像を学習させる技術を身につける必要があります。
本記事では、Deep Learningの研究開発で圧倒的なシェアを誇るフレームワーク「PyTorch」を用いて、ゼロからオリジナル画像分類モデル（畳み込みニューラルネットワーク: CNN）を作成し、学習を行うパイプラインの基礎を長文で徹底解説します。

1. PyTorchにおけるデータローダー（Dataset / DataLoader）の構築

PyTorchの最大の強みは「データの前処理とバッチ処理」を洗練されたオブジェクト指向で記述できる点です。自社工場で撮影した「良品・不良品」の画像などを学習させるには、まず独自の Dataset クラスを定義します。

import torch
import torchvision
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from PIL import Image
import os
class CustomImageDataset(Dataset):
    def __init__(self, img_dir, transform=None):
        # ディレクトリ内の画像ファイルパス一覧を取得
        self.img_dir = img_dir
        self.img_labels = [] 
        # 例：フォルダ名がクラス名であると仮定した読み込みロジックがここに入ります
        self.transform = transform
    def __len__(self):
        return len(self.img_labels)
    def __getitem__(self, idx):
        # 実際にデータが要求されたタイミングでディスクから画像を読み込む（メモリ節約）
        img_path = self.img_labels[idx][0]
        label = self.img_labels[idx][1]
        image = Image.open(img_path).convert("RGB")
        
        if self.transform:
            image = self.transform(image)
            
        return image, label
# データオーグメンテーションとTensor化の定義
transform = transforms.Compose([
    transforms.Resize((128, 128)),
    transforms.RandomHorizontalFlip(), # 水平反転でデータを水増し
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

2. CNN（畳み込みニューラルネットワーク）アーキテクチャの定義

続いて、画像の特徴を抽出する「畳み込み層（Conv2d）」と、画像のサイズをダウンサンプリングする「プーリング層（MaxPool2d）」を組み合わせたネットワークを nn.Module を継承して構築します。

import torch.nn as nn
import torch.nn.functional as F
class SimpleCNN(nn.Module):
    def __init__(self, num_classes=2):
        super(SimpleCNN, self).__init__()
        # 1層目の畳み込み: 入力チャネル3(RGB) -> 出力チャネル16
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, padding=1)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        
        # 2層目の畳み込み: 入力チャネル16 -> 出力チャネル32
        self.conv2 = nn.Conv2d(16, 32, 3, padding=1)
        
        # 全結合層 (128x128画像 -> 2回のプーリングで32x32になる)
        self.fc1 = nn.Linear(32 * 32 * 32, 128)
        self.fc2 = nn.Linear(128, num_classes)
    def forward(self, x):
        # Conv -> ReLU -> Pool
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        
        # 1次元のベクトルに平坦化
        x = x.view(-1, 32 * 32 * 32)
        
        # 全結合層
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

3. 学習ループ（Training Loop）の設計

モデルが定義できたら、最適化手法（Optimizer）と損失関数（Loss Function）を定義し、データを何度も繰り返し入力して重みを更新していく学習ループを回します。Deep Learningの計算は膨大なため、NVIDIAのGPU（CUDA）を使用するのが一般的です。

import torch.optim as optim
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = SimpleCNN(num_classes=2).to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
# 学習ループ (Epoch)
for epoch in range(10): 
    running_loss = 0.0
    for i, data in enumerate(dataloader, 0):
        inputs, labels = data
        inputs, labels = inputs.to(device), labels.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(dataloader):.4f}')
print("学習完了！")

まとめ：自作AIをローカル環境で本格的に動かすために

PyTorchの美しい設計思想により、複雑なニューラルネットワークも非常に直感的に記述できることがお分かりいただけたかと思います。しかし、画像認識モデルの学習には多大なGPUリソース（VRAM等）が必要です。
クラウドを利用する手もありますが、継続的に自前のデータセットを学習させる場合は、AI学習用のNVIDIA製グラフィックボードを搭載したPCを用意しておくのが結局コストパフォーマンスが良くなります。

リンク

PyTorchの実装方法について、より実践的で幅広いアーキテクチャ（ResNet、GAN、RNNなど）を網羅的に学習したい方は、以下の技術書がバイブルとして非常におすすめです。

リンク