Pytorch – DataLoader の使い方について解説

def _get_img_paths(img_dir):
    img_extensions = [".jpg", ".jpeg", ".png", ".bmp"]
    img_paths = [p for p in Path(img_dir).iterdir() if p.suffix in img_extensions]

    return img_paths


class ImageFolder(data.Dataset):
    def __init__(self, img_dir, transform=None):
        # 画像ファイルのパス一覧を取得する。
        self.img_paths = _get_img_paths(img_dir)
        self.transform = transform

    def __getitem__(self, index):
        # インデックスに対応するファイルパスを取得する。
        path = self.img_paths[index]
        # 画像を読み込む。
        img = Image.open(path)
        # Transforms で変換する。
        if transform is not None:
            img = self.transform(img)

        return img

    def __len__(self):
        # ディレクトリ内の画像枚数を返す。
        return len(self.img_paths)

transform = transforms.Compose(
    [transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor()]
)

dataset = ImageFolder("data", transform)

# データを取得する。
for i in range(len(dataset)):
    img = dataset[i]
    print(type(img), img.shape)

import numpy as np
from torch.utils import data as data


class MyDataset(data.Dataset):
    def __init__(self, n):
        self.data = np.random.rand(n, 100, 100, 3)
        self.labels = np.arange(n)

    def __getitem__(self, index):
        return self.data[index], self.labels[index]

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


dataset = MyDataset(10)

print("batchsize=None")
dataloader = data.DataLoader(dataset, batch_size=None)
X, y = next(iter(dataloader))
print(X.shape, y)

print("batchsize=int")
dataloader = data.DataLoader(dataset, batch_size=3)
X, y = next(iter(dataloader))
print(X.shape, y.shape)

print("shuffle=False")
dataloader = data.DataLoader(dataset, batch_size=5)
for X, y in dataloader:
    print(y, end=" ")

print("\nshuffle=True")
dataloader = data.DataLoader(dataset, batch_size=5, shuffle=True)
for X, y in dataloader:
    print(y, end=" ")

class SequentialSampler(data.Sampler):
    def __init__(self, dataset):
        self.dataset = dataset

    def __iter__(self):
        return iter(range(len(self.dataset)))


sampler = SequentialSampler(dataset)
for i in sampler:  # iterate すると、読み込むサンプルのキーを返す
    print(i, end=" ")

class BatchSampler(data.Sampler):
    def __init__(self, sampler, batch_size):
        self.sampler = sampler
        self.batch_size = batch_size

    def __iter__(self):
        batch = []
        for idx in self.sampler:
            batch.append(idx)
            if len(batch) == self.batch_size:
                yield batch
                batch = []
        if len(batch) > 0:
            yield batch


batch_sampler = BatchSampler(sampler, 3)
for batch_indices in batch_sampler:
    print(batch_indices)

import collections
import re

import numpy as np
import torch

np_str_obj_array_pattern = re.compile(r"[SaUO]")

default_collate_err_msg_format = (
    "default_collate: batch must contain tensors, numpy arrays, numbers, "
    "dicts or lists; found {}"
)


def default_collate(batch):
    elem = batch[0]
    elem_type = type(elem)

    print(f"elem_type.__name__: {elem_type.__name__}")

    if isinstance(elem, torch.Tensor):
        # Tensor の場合はそのまま結合
        return torch.stack(batch, 0)
    elif (
        elem_type.__module__ == "numpy"
        and elem_type.__name__ != "str_"
        and elem_type.__name__ != "string_"
    ):
        if elem_type.__name__ == "ndarray" or elem_type.__name__ == "memmap":
            # numpy 配列の場合
            if np_str_obj_array_pattern.search(elem.dtype.str) is not None:
                raise TypeError(default_collate_err_msg_format.format(elem.dtype))

            # スカラー以外の場合、テンソルに変換して、default_collate() にもう一度渡す。
            return default_collate([torch.as_tensor(b) for b in batch])
        elif elem.shape == ():
            # スカラーの場合、そのままテンソルにする。
            return torch.as_tensor(batch)
    elif isinstance(elem, float):
        # float の場合、そのままテンソルにする。
        return torch.tensor(batch, dtype=torch.float64)
    elif isinstance(elem, int):
        # int の場合、そのままテンソルにする。
        return torch.tensor(batch)
    elif isinstance(elem, (str, bytes)):
        # str, bytes の場合、そのまま返す。
        return batch
    elif isinstance(elem, collections.abc.Mapping):
        # Mapping の場合 (dict など)、キーごとにテンソル化する。
        return {key: default_collate([d[key] for d in batch]) for key in elem}
    elif isinstance(elem, tuple) and hasattr(elem, "_fields"):
        # namedtuple の場合
        return elem_type(*(default_collate(samples) for samples in zip(*batch)))
    elif isinstance(elem, collections.abc.Sequence):
        # sequence の場合
        it = iter(batch)
        elem_size = len(next(it))
        if not all(len(elem) == elem_size for elem in it):
            raise RuntimeError("each element in list of batch should be of equal size")
        transposed = zip(*batch)
        return [default_collate(samples) for samples in transposed]

    raise TypeError(default_collate_err_msg_format.format(elem_type))

# 各要素が (10, 10) の numpy 配列
batch = [np.random.rand(10, 10) for i in range(5)]
output = default_collate(batch)
print(f"output.shape: {output.shape}, output.dtype: {output.dtype}")

# 各要素が () の numpy 配列 (スカラー)
batch = [np.array(5) for i in range(5)]
output = default_collate(batch)
print(f"output.shape: {output.shape}, output.dtype: {output.dtype}")

# 各要素が int
batch = [1 for i in range(5)]
output = default_collate(batch)
print(f"output.shape: {output.shape}, output.dtype: {output.dtype}")

# 各要素が () の numpy 配列 (スカラー)
batch = [1.1 for i in range(5)]
output = default_collate(batch)
print(f"output.shape: {output.shape}, output.dtype: {output.dtype}")

# 各要素が str
batch = ["hoge" for i in range(5)]
output = default_collate(batch)
print(output)

# 各要素が str
batch = [{"A": np.random.rand(10, 10), "B": "hoge", "C": 1} for i in range(5)]
output = default_collate(batch)

for key, value in output.items():
    if isinstance(value, list):
        print(f"key: {key}, len(value): {len(value)}, type(value[0]): {type(value[0])}")
    else:
        print(f"key: {key}, value.shape: {value.shape}, value.dtype: {value.dtype}")

from collections import namedtuple

Sample = namedtuple("sample", ["A", "B", "C"])

# 各要素が str
batch = [Sample(np.random.rand(10, 10), "hoge", 1) for i in range(5)]
output = default_collate(batch)

for key, value in output._asdict().items():
    if isinstance(value, torch.Tensor):
        print(f"key: {key}, value.shape: {value.shape}, value.dtype: {value.dtype}")
    else:
        print(f"key: {key}, len(value): {len(value)}, type(value[0]): {type(value[0])}")

# 各要素が list
batch = [["hoge", 2, np.random.rand(2, 2)] for i in range(5)]
output = default_collate(batch)
for value in output:
    if isinstance(value, torch.Tensor):
        print(f"value.shape: {value.shape}, value.dtype: {value.dtype}")
    else:
        print(f"len(value): {len(value)}, type(value[0]): {type(value[0])}")

# 各要素が tuple
batch = [("hoge", 2, np.random.rand(2, 2)) for i in range(5)]
output = default_collate(batch)
for value in output:
    if isinstance(value, torch.Tensor):
        print(f"value.shape: {value.shape}, value.dtype: {value.dtype}")
    else:
        print(f"len(value): {len(value)}, type(value[0]): {type(value[0])}")

print(f"number of samples: {len(dataset)}")

dataloader = data.DataLoader(dataset, batch_size=128, drop_last=False)
print("drop_last=False")
for X, y in dataloader:
    print(X.shape, y.shape)

dataloader = data.DataLoader(dataset, batch_size=128, drop_last=True)
print("drop_last=True")
for X, y in dataloader:
    print(X.shape, y.shape)