画像読込時のカーネルクラッシュ解決方法（カスタムデータセット）

大きなデータセットを扱っていて、jupyterが頻繁にカーネルクラッシュを起こしてしまい悩んでいたのがカスタムデータセットを作成することで解決できたのでメモ書きします。

# 悩み
jupyterカーネルを使っているときに表示されるこちら。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3830602/8200b0e8-7875-42d5-b555-4b8a0c77188f.png)

画像のような大きなデータセットを扱っているとよく出会います。私は画像を読み込む際と、訓練と検証データに分ける`train_test_split`を使った際にクラッシュしました。

原因は**メモリの圧迫**です。

手軽な解決策としては、使う予定のない変数を以下のようにガベージコレクションに入れることです。
```python
del image_tensors
import gc
gc.collect()
```

他にもfloat32をfloat16にするといった方法もあると思いますが、これを行った際に、勾配が0.0になってしまいました。

ちなみにクラッシュ直前のメモリの使用状況を見てみると、
```python
import psutil
print(f"使用中メモリ: {psutil.virtual_memory().used / (1024 ** 3):.2f} GB")
```
**使用中メモリ: 13.97 GB**

約14GBも使っていました。
どうにかできないかと思ってたどり着いたのが、カスタムデータセットを使っての画像の読み込みです。どのような読み込み方法をしたかについて述べていきます。

# 一般的な画像の読み込み
今回はHAM10000という皮膚がんの多クラス分類のデータセットを使っていきます。
このデータセットには、10015枚の画像とメタデータのcsvファイルがあります。

ここで一般的な画像の読み込みは、
```python
image_path = "./data/images"

# 画像の前処理
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
])

image_files = [f for f in os.listdir(image_path) if f.endswith(('.jpg'))]

image_tensors = []

for file in image_files:
    img = accimage.Image(os.path.join(image_path, file)) # PILで画像を開く
    img = transform(img) # 前処理を適用
    image_tensors.append(img)

image_tensors = torch.stack(image_tensors)

print(f"Loaded {len(image_tensors)} images as tensors.")
```

この後、訓練データと検証データに分けるのに`train_test_split`を使うので、`np.transpose`で`numpy.darray`に変えます。なお、Tensorとnumpyの画像の表示形式が違うので注意です。
```python
# PyTorchの CHWフォーマットになっているので、NumPy形式の NHWCフォーマットに変換する
image_tensors = image_tensors.numpy()  # Tensor → NumPy
image_tensors_numpy = np.transpose(image_tensors, (0, 2, 3, 1))  # CHW → HWC (Num_samples, Height, Width, Channels)
print(image_tensors_numpy.shape)
```

この後データセットを`train_test_split`で分割していきます。stratify＝labelsを指定していますが、これはクラスごとで均等に分けるための引数です。今回のデータセットではクラス不均衡があったので指定しています。
この後、データセットとデータローダーを作成することが通常の流れだと思います。
```python
from sklearn.model_selection import train_test_split
train_data, val_data, train_labels, val_labels = train_test_split(
    image_tensors_numpy, labels, test_size=0.3, stratify=labels, random_state=42
)
print("train_data shape:", train_data.shape)
print("train_labels shape:", train_labels.shape)
print("val_data shape:", val_data.shape)
print("val_labels shape:", val_labels.shape)
```
**私はここでよくカーネルクラッシュを起こしました。**
単純にカーネルクラッシュを防ぐだけなら、データ量を減らせばよいだけですが、クラス不均衡を持ったデータなのでどうやって減らすかに一工夫加える必要があります。
私は、今回インデックスで先に訓練と検証に分けることで、それを解決しました。

# 解決法：カスタムデータセットを使用する
今回メタデータとして、下記のようなデータフレームがあるとします。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3830602/79a10861-90cd-4257-b841-4a8d7dfe0b92.png)
この情報を使ってデータを分割していきます。

## ラベルのマッピング
まずは、ラベルがカテゴリになっているので数値に変えていきます。
```python
labels = metadata_labels['dx'].map({"akiec":0, "bcc":1, "bkl":2, "df":3, "mel":4, "nv":5, "vasc":6})
labels
```
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3830602/6df9a35e-4a7f-48cb-90de-8eec69d05c35.png)
## インデックスを分割
先ほど、マッピングしたラベルをnumpy形式で`labels`に保存します。
```python
# ラベルを取得
labels = labels.values
```
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3830602/442102cd-089b-43ab-ac11-e5af8e91766b.png)
次に`train_test_split`を使って実際のデータではなく、**ラベルのインデックス**を分割します。
```python
# インデックスを分割
from sklearn.model_selection import train_test_split
train_idx, val_idx = train_test_split(
    np.arange(len(image_id)),
    test_size = 0.3,
    stratify=labels, # クラスごとに均等に分割
    random_state = 42
)
```

## カスタムデータセットの定義
ここでデータセットを2つ定義します。
1. ImageDataset ⇒ 画像とラベルを保存するデータセット
1. CustomSubset ⇒ 分割と変換のためのデータセット

まず、画像に施したい変換を定義します。
```python
train_transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.RandomHorizontalFlip(),
    transforms.RandomRotation(10),
    transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),
    transforms.RandomResizedCrop(224, scale=(0.8, 1.0)),
    transforms.ToTensor(), # PIL 画像を (C, H, W) の torch.Tensor に変換
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

val_transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
```
### カスタムデータセット① ImageDataset
一つ目のデータセットは、画像とラベルを取得してデータセットに保存するものです。
```python
# ImageDataset クラスを定義
from torch.utils.data import Dataset, DataLoader

image_path = "./data/images"

class ImageDataset(Dataset):
    def __init__(self, image_id, labels):
        self.image_id = [f if f.endswith('.jpg') else f + ".jpg" for f in image_id] # 拡張子を追加
        self.labels = labels

    def __len__(self):
        return len(self.image_id)
    
    def __getitem__(self, idx):
        # 画像を指定するディレクトリを作成
        img_path = os.path.join(image_path, self.image_id[idx]) # self.image_id[idx]は画像ファイル名
        
        # ファイルが存在しない場合のエラー
        if not os.path.exists(img_path):
            raise FileNotFoundError(f"File not found: {img_path}")
        
        # ディレクトリで画像を指定しPIL形式で開き、RGB形式に変換する
        img = Image.open(img_path).convert("RGB")
        # ラベルをTensorの整数に変換する
        label  = torch.tensor(self.labels[idx], dtype=torch.long)
        return img, label # 画像とラベルを返す
```

### カスタムデータセット② CustomSubset
二つ目のカスタムデータセットは、データセットに保存した画像とラベルを取り出し、訓練用、検証用の変換を施します。ここでは、先ほどtrain_test_splitで作成したtrain_idxとval_idxを使用します。
```python
# CustomSubsetクラス
class CustomSubset(Dataset):
    def __init__(self, dataset, indices, transform=None):
        self.dataset = dataset
        self.indices = indices
        self.transform = transform
    
    def __len__(self):
        return len(self.indices)
    
    def __getitem__(self, idx):
        img, label = self.dataset[self.indices[idx]] # データセットからインデックス情報を用いて取得
        # transformが指定されれば実施（デフォルトはNone）
        if self.transform:
            img = self.transform(img)
        
        return img, label # 画像とラベルを返す
```
## データセットを使い、画像の読込と分割を行う
定義した2つのカスタムデータセットを使用します。
```python
# 画像とラベルのデータセットを作成
dataset = ImageDataset(image_id, labels)
```

```python
# インデックスを使用してデータセットの訓練用と検証用に分割
train_dataset = CustomSubset(dataset, train_idx, transform=train_transform)
val_dataset = CustomSubset(dataset, val_idx, transform=val_transform)
```

## データローダーを作成
それでは、データローダーを使い、バッチサイズ32に分けていきます。
```python
# Dataloaderの作成
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False)
```

データセットの中身を見ていきましょう。
```python
for imgs, labels in train_loader:
    print(f"Train Batch shape: {imgs.shape}, Labels: {labels}")
    break

for imgs, labels in val_loader:
    print(f"Val Batch shape: {imgs.shape}, Labels: {labels}")
    break
```
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3830602/4b39e0dd-18ef-4ead-87aa-6bcdc356adb2.png)

最後に使用メモリを確認します。
```python
import psutil
print(f"使用中メモリ: {psutil.virtual_memory().used / (1024 ** 3):.2f} GB")
```
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3830602/342e009b-9ca0-40f1-b414-b0986352fadb.png)

13.97 GBから4.57GBに変わりました！
つまり、**使用メモリを9.4GB減らせる**ことができました！
****
以上です。読んでいただきありがとうございました。