【必見】PyTorchで作る最強画像データセット！前処理＆拡張の全手順公開

2025年3月2日

どうもニコイチです。第4回の前回までで、zipファイルから画像データを解凍し、フォルダ構造や画像数、重複の有無を確認しました。
しかし、単に画像がフォルダに保存されているだけでは、モデルは学習できません。
今回のデータセット作成では、画像データとそのラベルをペアにした状態にし、さらにデータの前処理や水増し（Data Augmentation）を施すことで、効率的な学習を可能にします。

データセット作成の基本フロー

機械学習では、以下の3点セットでデータセットを作成します。

transforms
→ 画像の前処理やデータ拡張（例：リサイズ、テンソル化、正規化、水平方向の反転など）を定義します。
Dataset
→ 前処理を適用した上で、画像とそのラベルの組（ペア）を返すモジュールです。
今回は、フォルダ構造が「クラスごと」に整理されているため、datasets.ImageFolder()を使用します。
DataLoader
→ Datasetからデータをバッチ単位で取り出すためのモジュールです。
バッチサイズやシャッフルの有無、並列処理数などを指定できます。

全体の流れは、
transforms → Dataset → DataLoader
となっており、これらを組み合わせることで、モデル学習時に必要なデータが自動的に供給される仕組みが完成します。

transformsの定義

まずは、torchvisionのtransformsモジュールを用いて、画像の前処理とデータ拡張を定義します。
ここでは、以下の4つの処理を行います。

画像サイズの統一
→ transforms.Resize(256)で画像のサイズを256に統一します。
テンソル化
→ transforms.ToTensor()で画像データをテンソル（数値データ）に変換します。
正規化
→ transforms.Normalize()で、Imagenetの平均 [0.485, 0.456, 0.406] と標準偏差 [0.229, 0.224, 0.225] を用いて正規化します。
水平方向へのランダム反転（Data Augmentation）
→ 学習データ（train）のみ、transforms.RandomHorizontalFlip(p=0.5)でランダムに反転処理を行います。

学習用（train）と検証用（val）では、データ拡張の処理が異なるため、下記のように分けて定義します。

from torchvision import transforms

data_transforms = {
    'train': transforms.Compose([
        transforms.Resize(256),
        transforms.RandomHorizontalFlip(p=0.5),  # 学習時のみ反転
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
    ]),
    'val': transforms.Compose([
        transforms.Resize(256),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
    ]),
}

# 定義内容を確認してみましょう
print(data_transforms)

このように、data_transformsに学習用と検証用の前処理がまとめられています。

Datasetの作成

次に、前回解凍した画像データを使って、Datasetを作成します。
フォルダ構造がクラス（ラベル）ごとに整理されている場合、PyTorchのdatasets.ImageFolderを利用することで、各画像に自動的にラベル（例えば、’ng’が0、’ok’が1）が割り振られます。

以下のコードで、学習用と検証用のデータセットを作成します。

from torchvision import datasets

# datasets.ImageFolder()を用いてDatasetを作成
image_datasets = {
    'train': datasets.ImageFolder('./image_data/train', data_transforms['train']),
    'val': datasets.ImageFolder('./image_data/val', data_transforms['val'])
}

# 作成したDatasetの概要を確認
print(image_datasets['train'])

実行結果には、読み込んだデータの件数やルートパス、適用したtransformsが表示され、例えば「Number of datapoints: 120」などと出力されます。
また、image_datasets['train'].class_to_idxで、クラスごとのラベル割り振りも確認できます。

DataLoaderの作成

最後に、作成したDatasetからデータをバッチ単位で取り出すために、DataLoaderを作成します。
ここでは、学習用はデータをシャッフルしながらバッチサイズ4で、検証用はシャッフルなしでバッチサイズ4、さらにdrop_last=Trueで余ったデータをスキップする設定とします。

import torch

image_dataloaders = {
    'train': torch.utils.data.DataLoader(image_datasets['train'], batch_size=4,
                                           shuffle=True, num_workers=0, drop_last=True),
    'val': torch.utils.data.DataLoader(image_datasets['val'], batch_size=4,
                                         shuffle=False, num_workers=0, drop_last=True),
}

# DataLoaderの型を確認してみましょう
print(type(image_dataloaders['train']))

これで、DataLoaderを通して、モデル学習に必要なデータがバッチ単位で供給されるようになりました。
実際にデータが正しく取り出せるか、以下のようにfor文で確認してみます。

# DataLoaderから1バッチ分のデータを取り出して表示
for i, (inputs, labels) in enumerate(image_dataloaders['train']):
    print("入力データのテンソルサイズ：", inputs.size())
    print("ラベル：", labels)
    if i == 0:  # 1バッチだけ表示
        break

課題

以下の選択肢の中から、データセット作成の流れについて誤っているものを1つ選んでください。

PCスペックが高い場合はバッチサイズを大きくし、低い場合は小さくし、PCの性能に合わせて学習させることができる。
dataloadersがない場合はモデルの学習は不可能である。
dataloadersの型はtorch.utils.data.dataloader.DataLoaderである。

【正解】 2. dataloadersがない場合はモデルの学習は不可能である

【解説】
実際には、DataLoaderを使用しなくてもDatasetから直接データを取り出して学習することは可能ですが、バッチサイズやシャッフルなどの便利な機能を実装するためにはDataLoaderが必要です。
よって、選択肢2は「誤っている」とされています。

まとめ

今回の記事では、画像分類モデルの学習に必要なデータセット作成の流れを、以下の3つのステップで実践しました。

transforms
画像の前処理（リサイズ、テンソル化、正規化）と、学習用データに対する水平方向のランダム反転によるData Augmentationを定義。
Dataset
datasets.ImageFolderを利用し、フォルダ構造から画像と自動的にラベルのペアを作成。
DataLoader
Datasetからバッチ単位でデータを取り出し、学習に適した形にデータを整形。

これらのステップを正しく実行することで、モデルが効率的に学習できるデータセットが完成します。次回は、このデータセットを用いて実際にモデルの選定と学習に挑戦していきます！