PyTorchのDataLoaderにおけるCPUメモリに乗り切らないデータ(h5)のハンドリング

前提とやりたいこと

• h5ファイルの1つのdatasetに特徴量がドカッとまとめられている(N, features)
• 特徴量は既にシャッフルされており、dataloaderでシャッフルし直す必要はない
• N個のデータセットはCPUメモリに一度に載せることができないぐらい大きく、一度にn個しか載せられない(N = O(1)*n ~ O(10)*n)
• torch.utils.data.Datasetを継承したカスタムデータセットクラスを作成することでindexを指定して1つずつh5ファイルから読み出すこともできるが信じられないぐらい遅い
• ので、一度にn個をドカっと読み込みたい

解決法

• カスタムデータセットクラスの__getitem__メソッドを変更する
• DataLoaderのbatch_sizeを1に設定し、collate_fnをカスタマイズして、__getitem__メソッドで取得したデータセットをそのまま返すようにする
• 一度に読み込む量はpreload_factorで設定する
• 学習時にbatch_sizeごとに小分けして取り出す

import h5py
import torch
from torch.utils.data import Dataset

class H5Dataset(Dataset):
    def __init__(self, h5_file, preload_factor, transform=None):
        self.h5_file = h5_file
        self.preload_factor = preload_factor
        self.transform = transform

    def __len__(self):
        with h5py.File(self.h5_file, 'r') as f:
            return (len(f['features']) - 1) // self.preload_factor + 1

    def __getitem__(self, idx):
        with h5py.File(self.h5_file, 'r') as f:
            start = idx * self.preload_factor
            end = start + self.preload_factor
            features = f['features'][start:end]
            labels = f['labels'][start:end]


        if self.transform:
            features = self.transform(features)

        return torch.tensor(features), torch.tensor(labels)

def collate_fn(batch):
    return batch[0]

h5_file = 'your_h5_file_path.h5'
batch_size = 1000
preload_factor = 1000000

dataset = H5Dataset(h5_file, preload_factor=preload_factor)
dataloader = DataLoader(dataset, batch_size=1, shuffle=False, collate_fn=collate_fn)

for batch in dataloader:
    # バッチの処理
    for i in range(0, batch.shape[0], batch_size):
        mini_batch = batch[i:i+batch_size]
        # ここで深層学習モデルの学習を行う

注意

• h5ファイルの時点でシャッフルされている必要がある（されていない場合、通常の方法で作成したdataloaderでshuffle=Trueした場合と、この方法でshuffle=Trueした場合とで挙動が異なってしまう為）