【Pythonお悩み解決】データセットがメモリにのらないときの対処法

この記事は， Pythonを利用して研究を行なっていく中で私がつまずいてしまったポイントをまとめていくものです。同じような状況で苦しんでいる方々の参考になれば嬉しいです。Pythonつまずきポイント集の目次は以下のページをご覧ください。

【超初心者お悩み解決】Pythonつまずきポイント記事まとめページ この記事は，Pythonを利用して研究を行なっていく中で私がつまずいてしまったポイントをまとめていくものです。同じような状況で苦しんで...

環境

現象

対処法

少しトリッキーですが，datasetに渡すのは「予めバッチサイズに分割されたデータが格納されたpath名のリスト」とします。こちらの記事（【超初心者向け】ド素人がPyTorchで自作データセットを作るまで。）でも説明しているように，普通pytorchデフォルトのdatasetを使っている場合はdatasetクラスのインスタンスを生成するために全データそのものを渡します。しかし，使用するデータが莫大な場合には，データセットを一旦全部読み込んでしまうため，メモリが大量消費されます。

この現象を回避するためには，毎回のイテレーションごとにデータセットを読み込む必要があります。そこで，pytorchにdatasetとして認識させるのは，データセットのpath名が格納されたリストだということにします。そうすれば，pytorchのdatasetとdataloaderが「どのバッチを選ぶのか」というインデックス処理をうまいことやってくれます。

実際のコードの一部はこんな感じです。

class LabelDataset(torch.utils.data.Dataset):
    '''
    data: N個のnumpyテンソルのpathが格納されているリスト
    label: N個のnumpyテンソルのpathが格納されているリスト
    transform: とりあえずテンソル化のみ
    '''
    def __init__(self, data, label, transform=None):
        self.transform = transform
        self.data = data
        self.data_num = len(data)
        self.label = label

    def __len__(self):
        return self.data_num

    def __getitem__(self, idx):
        out_data = self.transform(np.load(self.data[idx]))
        out_label = self.transform(np.load(self.label[idx]))
        return out_data, out_label

transform = transforms.Compose([torch.tensor])

# データセットのインスタンス化
train_dataset = my_dataset.LabelDataset(data=train_data, label=label, transform=transform)
train_batch_size = 256
train_loader = DataLoader(train_dataset, batch_size=train_batch_size, shuffle=True)

# 学習用のパラメータ設定
num_epochs = 100
learning_rate = 1e-5

# モデルのインスタンス化
model = CNN()
criterion = nn.functional.binary_cross_entropy
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate, weight_decay=1e-5)

###############################################################################
# 学習開始
model.train()
model.to(device)
loss_list = []
loss_all = []

for epoch in range(num_epochs):
    losses = []
    count = 0
    for x,t in train_loader:
        x = x.to(device).float()
        t = t.to(device).float()
        model.zero_grad()
        y = model(x)
        loss = criterion(y,t)
        loss.backward()
        optimizer.step()
        losses.append(loss.cpu().detach().numpy())
        loss_all.append(loss.cpu().detach().numpy())
        sys.stdout.write("\rBatch No. [{}/{}]".format(count+1, len(train_loader)))
        sys.stdout.flush()
        count += 1
        loss_all.append(losses)
    loss_list.append(np.average(losses))
    sys.stdout.write("\rEpoch [{}/{}], Loss: {:.4f}\n".format(epoch + 1, num_epochs, np.average(losses)))
print("-----学習終了-----")

ひとこと

深層学習のメモリ問題は深刻です…。データセットが確定しているのであれば，入力特徴量をバッチサイズに分割しておいて特定のディレクトリに保存しておけば，上手いことメモリを節約しながら学習を回すことが可能になります。