PyTorchのDataset,Dataloaderの学習時間が長くて困っている人へ

# はじめに
こんにちは！大学院生で機械学習の研究をしている者です。
先日、学習に使用していたコードをリファクタリングしていた際、ある部分を修正したところ、学習時間が大幅に短縮されました。
その話を研究室の同期にしたところ、彼も同じミスをしていたことがわかり、「これは他にも同じミスをしている人がいるのではないか」と思い、今回この記事を書くことにしました。

初投稿であり、私自身まだ機械学習歴が半年ほどの初心者ですので、もし間違いや改善点があればぜひ教えていただけると嬉しいです！

# 結論
先に結論から書くと`__getitem__`で画像を読み込む処理を行っていたことが原因でした。
データを読み込む処理を`__init__`で実行することによって、学習時間を短くすることができます。

# 実例
Datasetを自作する場合、以下の3つのメソッドを実装する必要があります。
(1) `__init__`:
データの読み込みや初期設定を行います。

(2) `__len__`:
データセットの総サンプル数を返します。

(3) `__getitem__`:
インデックスを指定して、データを1つ取得できるようにします。

以下に<strong>悪い例</strong>を載せてみます。
```Python
class CustomImageDataset(Dataset):
    def __init__(self, image_dir, transform=None):
        """
        カスタムデータセットの初期化
        Args:
            image_dir (str): 画像が格納されたディレクトリのパス
            transform (callable, optional): データ前処理やデータ拡張を行う関数
        """
        self.image_dir = image_dir
        self.transform = transform
        self.image_paths = []
        
        # ディレクトリ内の全ての画像ファイルパスを収集
        for root, _, files in os.walk(image_dir):
            for file in files:
                if file.endswith(('.png', '.jpg')):
                    self.image_paths.append(os.path.join(root, file))

    def __len__(self):
        """データセットの長さを返す"""
        return len(self.image_paths)

    def __getitem__(self, index):
        """
        データセットから特定のインデックスの画像データを取得
        Args:
            index (int): 取得したいデータのインデックス
        """
        # 画像ファイルを読み込む
        image = Image.open(self.image_paths[index])
        
        label = os.path.basename(os.path.dirname(image_path))
        
        if self.transform:
            image = self.transform(image)  # データ拡張を適用
        
        return image, label
```

このコードで問題なのは`__getitem__`の画像を読み込んでいる処理です。
```Python
        image = Image.open(self.image_paths[index])
```
この処理を`__init__`で実行し、`self.image = []`に読み込んだ画像を格納しておくことで学習時間を短くすることができます。
以下に修正例を載せておきます。

```Python
class CustomImageDataset(Dataset):
    def __init__(self, image_dir, transform=None):
        """
        カスタムデータセットの初期化
        Args:
            image_dir (str): 画像が格納されたディレクトリのパス
            transform (callable, optional): データ前処理やデータ拡張を行う関数
        """
        self.image_dir = image_dir
        self.transform = transform
        self.image_paths = []
        self.image = [] # 追加、読み込んだ画像を保存
        
        # ディレクトリ内の全ての画像ファイルパスを収集
        for root, _, files in os.walk(image_dir):
            for file in files:
                if file.endswith(('.png', '.jpg')):
                    self.image_paths.append(os.path.join(root, file))

        # 画像ファイルを読み込む
        for i in range(len(self.image_paths)):
            image = Image.open(self.image_paths[index])
            self.image.append(image)

    def __len__(self):
        """データセットの長さを返す"""
        return len(self.image_paths)

    def __getitem__(self, index):
        """
        データセットから特定のインデックスの画像データを取得
        Args:
            index (int): 取得したいデータのインデックス
        """
        image = self.image[index]
        label = os.path.basename(os.path.dirname(image_path))
        
        if self.transform:
            image = self.transform(image)  # データ拡張を適用
        
        return image, label
```

# なぜ短くなるのか
Dataloaderを使用して学習を行う際の簡単なサンプルを以下に載せます。
```Python
# 学習ループ
num_epochs = 100
for epoch in range(num_epochs):
    model.train()  # モデルを学習モードに設定
    running_loss = 0.0

    for inputs, labels in train_loader:
        # 入力データをモデルに流す
        outputs = model(inputs)
        loss = criterion(outputs, labels)  # 損失を計算

        # 勾配の初期化、逆伝播、パラメータ更新
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
```
`__getitem__`は以下の部分で実行されています。
```Python
    for inputs, labels in train_loader:
```
サンプルコードを見ればお分かりかと思いますが`__getitem__`内にデータの読み込み処理が入っている場合、同じデータをエポックごとに読み込んでいます。
そのため、データ数×（エポック数-1）回の読み込み時間を削減することができます。

# おまけ
## データ拡張について
データ拡張は`__getitem__`に書く必要があります。`__getitem__`はバッチごとに実行されるため、バッチごとに異なる拡張を行ったデータを学習に使用することができ、過学習を抑制できます。
逆に`__init__`に書いてしまうと意味がなくなってしまいます。

## ミスの原因を考察
1. 挙動がわかりづらい
`__getitem__`と直接コードに書かれているわけではないため、なかなか見つけづらいミスかなと感じました。

2. わかりやすいサイトがこの書き方
検索上位に表示されるサイトの中にもこの書き方をしているサイトがいくつかあったのでこれも原因じゃないかと思います。

3. ChatGPTもこの書き方
プロンプトを工夫したり、何回か実行したりすれば変わるのかもしれませんが、私が試しに生成させてみたときは`__getitem__`にデータの読み込み処理を書いていました。

4. PyTorchのチュートリアルもこの書き方？
PyTorchのチュートリアルもなぜか`__getitem__`にデータの読み込み処理を書いていました。使用していた関数の動きがわからないのでもしかしたら違うかもしれませんが、、、