PyTorch入門 第3回：Dataset、DataLoader、画像前処理とデータ拡張

yen_images/
  raw/
    1000_yen/
      img_0001.jpg
      img_0002.jpg
      ...
    10000_yen/
      img_0001.jpg
      img_0002.jpg
      ...

ローカル画像
  -> データ確認
  -> train / valid / test に分割
  -> Dataset: パスと __getitem__ を定義
  -> transforms: 前処理とデータ拡張
  -> Sampler / BatchSampler: index と順序を決定
  -> DataLoader worker: 画像を読み込む
  -> collate_fn: サンプルをbatchにまとめる
  -> モデルの訓練または評価

yen_images/
  split/
    train/
      1000_yen/
      10000_yen/
    valid/
      1000_yen/
      10000_yen/
    test/
      1000_yen/
      10000_yen/

from pathlib import Path
import random
import shutil


IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".bmp"}
CLASS_NAMES = ["1000_yen", "10000_yen"]


def collect_images(class_dir: Path) -> list[Path]:
    return sorted(
        path for path in class_dir.iterdir()
        if path.is_file() and path.suffix.lower() in IMAGE_EXTENSIONS
    )


def split_image_dataset(
    source_root: Path,
    output_root: Path,
    train_ratio: float = 0.8,
    valid_ratio: float = 0.1,
    seed: int = 42,
) -> None:
    if not 0 < train_ratio < 1:
        raise ValueError("train_ratio must be between 0 and 1")
    if not 0 <= valid_ratio < 1:
        raise ValueError("valid_ratio must be between 0 and 1")
    if train_ratio + valid_ratio >= 1:
        raise ValueError("train_ratio + valid_ratio must be smaller than 1")

    random_generator = random.Random(seed)

    for class_name in CLASS_NAMES:
        image_paths = collect_images(source_root / class_name)
        random_generator.shuffle(image_paths)

        train_end = int(len(image_paths) * train_ratio)
        valid_end = int(len(image_paths) * (train_ratio + valid_ratio))

        splits = {
            "train": image_paths[:train_end],
            "valid": image_paths[train_end:valid_end],
            "test": image_paths[valid_end:],
        }

        for split_name, paths in splits.items():
            target_dir = output_root / split_name / class_name
            target_dir.mkdir(parents=True, exist_ok=True)

            for source_path in paths:
                shutil.copy2(source_path, target_dir / source_path.name)

        print(
            f"{class_name}: "
            f"train={len(splits['train'])}, "
            f"valid={len(splits['valid'])}, "
            f"test={len(splits['test'])}"
        )
        # 出力例（画像枚数は環境によって変わります。ここでは各クラス50枚の場合）:
        # 1000_yen: train=40, valid=5, test=5
        # 10000_yen: train=40, valid=5, test=5


if __name__ == "__main__":
    split_image_dataset(
        source_root=Path("yen_images/raw"),
        output_root=Path("yen_images/split"),
    )

from torch.utils.data import Dataset


class MyDataset(Dataset):
    def __init__(self, samples, transform=None):
        self.samples = samples
        self.transform = transform

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, index):
        image, label = self.samples[index]
        if self.transform is not None:
            image = self.transform(image)
        return image, label

from pathlib import Path
from PIL import Image
from torch.utils.data import Dataset


IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".bmp"}
CLASS_NAMES = ["1000_yen", "10000_yen"]


class BanknoteDataset(Dataset):
    def __init__(self, root, transform=None):
        self.root = Path(root)
        self.transform = transform
        self.class_to_index = {name: index for index, name in enumerate(CLASS_NAMES)}
        self.samples = []  # (画像パス, ラベル) のタプルを集めるリスト

        for class_name in CLASS_NAMES:
            class_dir = self.root / class_name
            if not class_dir.exists():
                raise FileNotFoundError(f"class directory not found: {class_dir}")

            for path in sorted(class_dir.iterdir()):  # 実行環境によらず順序を固定する
                if path.is_file() and path.suffix.lower() in IMAGE_EXTENSIONS:
                    self.samples.append((path, self.class_to_index[class_name]))

        if not self.samples:
            raise RuntimeError(f"no image files found under: {self.root}")

    def __len__(self):
        return len(self.samples)  # DataLoaderやSamplerが参照するサンプル数

    def __getitem__(self, index):
        image_path, label = self.samples[index]  # indexからパスとラベルを取り出す
        with Image.open(image_path) as image:
            image = image.convert("RGB")  # グレースケールやRGBAが混ざっていても3チャンネルにそろえる

        if self.transform is not None:
            image = self.transform(image)  # ここでtransformを適用してからTensorを返す

        return image, label

import torch
from torch.utils.data import DataLoader
from torchvision.transforms import v2


loader_transform = v2.Compose([
    v2.ToImage(),
    v2.Resize((224, 224), antialias=True),
    v2.ToDtype(torch.float32, scale=True),
])

train_dataset = BanknoteDataset("yen_images/split/train", transform=loader_transform)

train_loader = DataLoader(
    train_dataset,
    batch_size=32,
    shuffle=True,
    num_workers=0,
    pin_memory=torch.cuda.is_available(),
    drop_last=True,
)

for images, labels in train_loader:
    print(images.shape)
    # torch.Size([32, 3, 224, 224])
    print(labels.shape)
    # torch.Size([32])
    break

def main():
    ...


if __name__ == "__main__":
    main()

DataLoader
  -> Sampler または BatchSampler が index を作る
  -> DataLoader iterator が index を worker に渡す
  -> DatasetFetcher が Dataset[index] を呼ぶ
  -> Dataset.__getitem__ で画像とラベルを読む
  -> transform を適用する
  -> collate_fn がサンプルのリストをbatchにする
  -> 必要に応じてpin memoryを使う
  -> batchを返す

images = torch.stack([sample[0] for sample in batch])
labels = torch.tensor([sample[1] for sample in batch])

def detection_collate(batch):
    images, targets = zip(*batch)
    return list(images), list(targets)

drop_last=False: iterations = ceil(N / batch_size)
drop_last=True : iterations = floor(N / batch_size)

drop_last=False: 8+8+8+8+8+8+8+8+8+8+7 = 11 iterations
drop_last=True : 8+8+8+8+8+8+8+8+8+8   = 10 iterations

import torch
from torchvision.transforms import v2


basic_transform = v2.Compose([
    v2.ToImage(),
    v2.Resize((224, 224), antialias=True),
    v2.ToDtype(torch.float32, scale=True),
])

dataset = BanknoteDataset("yen_images/split/train", transform=basic_transform)

image, label = dataset[0]
print(image.shape)
# torch.Size([3, 224, 224])
print(image.dtype)
# torch.float32
print(image.min().item(), image.max().item())
# 出力例（値は元画像の内容によって変わります）: 0.00392156862745098 0.9921568627450981

output[c] = (input[c] - mean[c]) / std[c]

R' = (R - mean_R) / std_R
G' = (G - mean_G) / std_G
B' = (B - mean_B) / std_B

import torch
from torch.utils.data import DataLoader
from torchvision.transforms import v2


stats_transform = v2.Compose([
    v2.ToImage(),
    v2.Resize((224, 224), antialias=True),
    v2.ToDtype(torch.float32, scale=True),
])


def compute_mean_std(dataset, batch_size=64):
    loader = DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=0)
    channel_sum = torch.zeros(3)          # チャンネルごとの画素値の合計
    channel_squared_sum = torch.zeros(3)  # チャンネルごとの画素値の2乗の合計
    pixel_count = 0                       # 1チャンネルあたりの画素数の合計

    for images, _ in loader:
        channel_sum += images.sum(dim=(0, 2, 3))                # batch, 高さ, 幅方向に合計する
        channel_squared_sum += (images ** 2).sum(dim=(0, 2, 3))
        pixel_count += images.shape[0] * images.shape[2] * images.shape[3]

    mean = channel_sum / pixel_count
    variance = (channel_squared_sum / pixel_count - mean ** 2).clamp_min(0)  # E[X^2] - E[X]^2。丸め誤差での負値を防ぐ
    std = variance.sqrt()
    return mean.tolist(), std.tolist()


stats_dataset = BanknoteDataset("yen_images/split/train", transform=stats_transform)
mean, std = compute_mean_std(stats_dataset)

print(mean)
# 出力例（値は使用する画像データによって変わります）: [0.5005499124526978, 0.4998324513435364, 0.5008472204208374]
print(std)
# 出力例（値は使用する画像データによって変わります）: [0.1450541615486145, 0.14825528860092163, 0.14333893358707428]

def denormalize(image, mean, std):
    mean = torch.as_tensor(mean, dtype=image.dtype, device=image.device)[:, None, None]
    std = torch.as_tensor(std, dtype=image.dtype, device=image.device)[:, None, None]
    return image * std + mean

from torchvision.transforms import v2


center_crop = v2.Compose([
    v2.ToImage(),
    v2.Resize(256, antialias=True),
    v2.CenterCrop((224, 224)),
])

random_crop = v2.RandomCrop(
    size=(224, 224),
    padding=16,
    pad_if_needed=True,
    fill=0,
    padding_mode="constant",
)

from torchvision.transforms import InterpolationMode


random_resized_crop = v2.RandomResizedCrop(
    size=(224, 224),
    scale=(0.7, 1.0),
    ratio=(0.9, 1.1),
    interpolation=InterpolationMode.BILINEAR,
    antialias=True,
)

元画像
  -> scaleから切り出し面積を決める
  -> ratioから幅と高さの比率を決める
  -> 元画像内で位置を決める
  -> 切り出す
  -> interpolationで固定サイズへResizeする

+-----------+-----------+
| top-left  | top-right |
|           |           |
+-----+ center +--------+
| bottom-left | bottom-right
+-------------+---------+

ten_crop_transform = v2.Compose([
    v2.ToImage(),
    v2.Resize(256, antialias=True),
    v2.TenCrop((224, 224)),
    v2.Lambda(lambda crops: torch.stack([
        v2.ToDtype(torch.float32, scale=True)(crop)
        for crop in crops
    ])),
])

horizontal_flip = v2.RandomHorizontalFlip(p=0.5)
vertical_flip = v2.RandomVerticalFlip(p=0.5)

rotation = v2.RandomRotation(
    degrees=(-10, 10),
    interpolation=InterpolationMode.BILINEAR,
    expand=False,
    fill=0,
)

pad = v2.Pad(
    padding=(8, 16, 8, 16),
    fill=0,
    padding_mode="constant",
)

color_jitter = v2.ColorJitter(
    brightness=0.2,
    contrast=0.2,
    saturation=0.1,
    hue=0.02,
)

gray_3ch = v2.Grayscale(num_output_channels=3)
random_gray = v2.RandomGrayscale(p=0.1)

affine = v2.RandomAffine(
    degrees=(-8, 8),
    translate=(0.05, 0.05),
    scale=(0.95, 1.05),
    shear=(-3, 3),
    interpolation=InterpolationMode.BILINEAR,
    fill=0,
)

入力座標
  -> centerを基準にする
  -> rotation
  -> scale
  -> shear
  -> translation
  -> interpolationで画素値を決める
  -> fillで空白を埋める

random_erasing = v2.RandomErasing(
    p=0.25,
    scale=(0.02, 0.15),
    ratio=(0.3, 3.3),
    value=0,
    inplace=False,
)

Tensor画像
  -> pで実行するか決める
  -> scaleで遮蔽面積を決める
  -> ratioで遮蔽領域の幅高さ比を決める
  -> 位置を決める
  -> valueまたはrandom値で埋める

import torch


class AddSaltPepperNoise:
    def __init__(
        self,
        noise_probability=0.03,
        apply_probability=0.5,
        salt_value=1.0,
        pepper_value=0.0,
    ):
        if not 0 <= noise_probability <= 1:
            raise ValueError("noise_probability must be in [0, 1]")
        if not 0 <= apply_probability <= 1:
            raise ValueError("apply_probability must be in [0, 1]")

        self.noise_probability = noise_probability
        self.apply_probability = apply_probability
        self.salt_value = salt_value
        self.pepper_value = pepper_value

    def __call__(self, image):
        if torch.rand((), device=image.device) >= self.apply_probability:  # 確率的にノイズを適用しない画像もつくる
            return image

        output = image.clone()  # 元のTensorを直接書き換えない
        mask = torch.rand(image.shape[-2:], device=image.device)  # 画素ごとに1つ乱数を割り当てる
        pepper = mask < self.noise_probability / 2       # 下側q/2はごま（黒）にする
        salt = mask > 1 - self.noise_probability / 2      # 上側q/2は塩（白）にする
        output[..., pepper] = self.pepper_value
        output[..., salt] = self.salt_value
        return output

transform = v2.Compose([
    v2.ToImage(),
    v2.Resize((224, 224), antialias=True),
    v2.ToDtype(torch.float32, scale=True),
])

画像読み込み
  -> 空間変換
  -> 色変換
  -> ToImage
  -> ToDtype(float32, scale=True)
  -> Normalize
  -> 必要ならTensor向けの遮蔽

choice = v2.RandomChoice([
    v2.RandomRotation(8),
    v2.ColorJitter(brightness=0.15),
])

apply = v2.RandomApply([
    v2.ColorJitter(brightness=0.2, contrast=0.2),
], p=0.5)

order = v2.RandomOrder([
    v2.RandomRotation(5),
    v2.ColorJitter(brightness=0.1),
])

import torch
from torchvision.transforms import InterpolationMode
from torchvision.transforms import v2


mean = [0.50, 0.50, 0.50]
std = [0.25, 0.25, 0.25]

train_transform = v2.Compose([
    v2.ToImage(),
    v2.RandomResizedCrop(
        (224, 224),
        scale=(0.75, 1.0),
        ratio=(0.9, 1.1),
        interpolation=InterpolationMode.BILINEAR,
        antialias=True,
    ),
    v2.RandomRotation(
        degrees=(-8, 8),
        interpolation=InterpolationMode.BILINEAR,
        fill=0,
    ),
    v2.ColorJitter(
        brightness=0.18,
        contrast=0.18,
        saturation=0.08,
        hue=0.01,
    ),
    v2.ToDtype(torch.float32, scale=True),
    AddSaltPepperNoise(noise_probability=0.02, apply_probability=0.25),
    v2.Normalize(mean=mean, std=std),
    v2.RandomErasing(p=0.2, scale=(0.02, 0.12), value=0),
])

valid_transform = v2.Compose([
    v2.ToImage(),
    v2.Resize(256, antialias=True),
    v2.CenterCrop((224, 224)),
    v2.ToDtype(torch.float32, scale=True),
    v2.Normalize(mean=mean, std=std),
])

import matplotlib.pyplot as plt


def show_samples(dataset, mean, std, count=6):
    figure, axes = plt.subplots(1, count, figsize=(count * 2.5, 2.5))
    for axis in axes:
        image, label = dataset[0]
        image = denormalize(image, mean, std).clamp(0, 1)
        axis.imshow(image.permute(1, 2, 0).numpy())
        axis.set_title(CLASS_NAMES[label])
        axis.axis("off")
    plt.tight_layout()
    plt.show()
    # 実行結果は画像のため誌面には貼りません。dataset[0]をcount回読み直した、
    # train_transformの拡張がかかったcount枚の画像が横一列に並び、各画像の上に
    # クラス名（千円札または一万円札）がタイトルとして表示されます。


preview_dataset = BanknoteDataset("yen_images/split/train", transform=train_transform)
show_samples(preview_dataset, mean, std)

from pathlib import Path
import random
import numpy as np
import torch
from torch import nn
from torch.utils.data import DataLoader


def set_seed(seed=42):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


class SmallCNN(nn.Module):
    def __init__(self, num_classes=2):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 32, kernel_size=3, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d((1, 1)),
        )
        self.classifier = nn.Linear(128, num_classes)

    def forward(self, x):
        x = self.features(x)
        x = torch.flatten(x, 1)
        return self.classifier(x)


def train_one_epoch(model, loader, criterion, optimizer, device, pin_memory):
    model.train()  # BatchNormやDropoutを訓練モードにする
    running_loss = 0.0
    correct = 0
    total = 0

    for images, labels in loader:
        images = images.to(device, non_blocking=pin_memory)
        labels = labels.to(device, non_blocking=pin_memory)

        logits = model(images)
        loss = criterion(logits, labels)

        optimizer.zero_grad(set_to_none=True)  # 前のbatchの勾配をリセットする
        loss.backward()
        optimizer.step()

        batch_size = labels.size(0)
        running_loss += loss.item() * batch_size  # batch平均をサンプル数で重み付けして合計する
        correct += (logits.argmax(dim=1) == labels).sum().item()
        total += batch_size

    return running_loss / total, correct / total


def evaluate(model, loader, criterion, device, pin_memory):
    model.eval()  # BatchNormやDropoutを評価モードにする
    running_loss = 0.0
    correct = 0
    total = 0

    with torch.inference_mode():  # 評価だけなので勾配計算を止める
        for images, labels in loader:
            images = images.to(device, non_blocking=pin_memory)
            labels = labels.to(device, non_blocking=pin_memory)

            logits = model(images)
            loss = criterion(logits, labels)

            batch_size = labels.size(0)
            running_loss += loss.item() * batch_size
            correct += (logits.argmax(dim=1) == labels).sum().item()
            total += batch_size

    return running_loss / total, correct / total


def main():
    set_seed(42)

    data_root = Path("yen_images/split")
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    num_workers = 0
    pin_memory = device.type == "cuda"

    train_dataset = BanknoteDataset(data_root / "train", transform=train_transform)
    valid_dataset = BanknoteDataset(data_root / "valid", transform=valid_transform)

    train_loader = DataLoader(
        train_dataset,
        batch_size=32,
        shuffle=True,
        num_workers=num_workers,
        pin_memory=pin_memory,
        drop_last=True,
        persistent_workers=num_workers > 0,
    )
    valid_loader = DataLoader(
        valid_dataset,
        batch_size=64,
        shuffle=False,
        num_workers=num_workers,
        pin_memory=pin_memory,
        drop_last=False,
        persistent_workers=num_workers > 0,
    )

    model = SmallCNN(num_classes=2).to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.5)

    best_valid_acc = 0.0
    max_epochs = 10

    for epoch in range(1, max_epochs + 1):
        train_loss, train_acc = train_one_epoch(
            model, train_loader, criterion, optimizer, device, pin_memory
        )
        valid_loss, valid_acc = evaluate(
            model, valid_loader, criterion, device, pin_memory
        )
        scheduler.step()  # StepLRの規定に従って学習率を更新する

        if valid_acc > best_valid_acc:  # 検証accuracyが更新されたときだけ保存する
            best_valid_acc = valid_acc
            torch.save(model.state_dict(), "best_banknote_cnn.pt")

        print(
            f"epoch {epoch:02d}: "
            f"train_loss={train_loss:.4f}, train_acc={train_acc:.3f}, "
            f"valid_loss={valid_loss:.4f}, valid_acc={valid_acc:.3f}"
        )
        # 出力例（ランダム生成した仮の画像で10 epoch実行した場合。実際の値は使用する画像データや乱数によって変わります）:
        # epoch 01: train_loss=0.6996, train_acc=0.484, valid_loss=0.6936, valid_acc=0.500
        # epoch 02: train_loss=0.7010, train_acc=0.484, valid_loss=0.6933, valid_acc=0.500
        # epoch 03: train_loss=0.6915, train_acc=0.547, valid_loss=0.6932, valid_acc=0.500
        # epoch 04: train_loss=0.7073, train_acc=0.500, valid_loss=0.6939, valid_acc=0.500
        # epoch 05: train_loss=0.7140, train_acc=0.500, valid_loss=0.6958, valid_acc=0.500
        # epoch 06: train_loss=0.7241, train_acc=0.453, valid_loss=0.6985, valid_acc=0.500
        # epoch 07: train_loss=0.6917, train_acc=0.500, valid_loss=0.7001, valid_acc=0.500
        # epoch 08: train_loss=0.6904, train_acc=0.531, valid_loss=0.7021, valid_acc=0.500
        # epoch 09: train_loss=0.7017, train_acc=0.469, valid_loss=0.7040, valid_acc=0.500
        # epoch 10: train_loss=0.6935, train_acc=0.500, valid_loss=0.7027, valid_acc=0.500


if __name__ == "__main__":
    main()

from PIL import Image


DISPLAY_NAMES = {
    0: "千円札",
    1: "一万円札",
}


def predict_image(image_path, model, transform, device):
    model.eval()
    with Image.open(image_path) as image:
        image = image.convert("RGB")

    x = transform(image).unsqueeze(0).to(device)

    with torch.inference_mode():
        logits = model(x)
        predicted_label = logits.argmax(dim=1).item()

    return DISPLAY_NAMES[predicted_label]

import torch
from torchvision import tv_tensors
from torchvision.transforms import v2


height, width = 480, 640
image = tv_tensors.Image(torch.randint(0, 256, (3, height, width), dtype=torch.uint8))
boxes = tv_tensors.BoundingBoxes(
    torch.tensor([[120, 80, 420, 360]], dtype=torch.float32),
    format="XYXY",
    canvas_size=(height, width),
)
mask = tv_tensors.Mask(torch.zeros((height, width), dtype=torch.uint8))

joint_transform = v2.Compose([
    v2.RandomHorizontalFlip(p=0.5),
    v2.RandomResizedCrop((320, 320), antialias=True),
    v2.ToDtype(torch.float32, scale=True),
])

image, boxes, mask = joint_transform(image, boxes, mask)