grid-searchによりpythonスクリプトが繰り返し実行される度に、RAM,GPUのメモリが解放されないまま増え続ける

Question

grid-searchによりpythonスクリプトが繰り返し実行される度に、RAM,GPUのメモリが増え続ける

はじめまして、当方teratailでの質問が始めてなので、質問の仕方など粗が目立ちますがご了承下さい。

発生している問題・エラーメッセージ

GCNというグラフを解析するニューラルネットのモデルがあるのですが、そのGCNというモデルのハイパーパラメータがgrid-searchにより最適化されるまで、モデルの学習とテストを繰り返すようなデモプログラムを作りました。

私の環境ではRAM,GPUのメモリ共に少ないため、モデルの学習が終わる度にメモリの使用量がリセットされるように動作してほしいです。

しかし実際に動かしてみると、モデルの学習が終わっても、メモリがほとんど解放されないまま次のハイパーパラメータでモデルの学習が再開されてしまいます。
そのため、数回モデルの学習を繰り返すとOut of Memoryになります。

PyTorchだけでなく、PyTorch_geometric、ogb、hydraなどthird-partyのライブラリが多く読みにくいかもしれませんが、
どんなアドバイスでも大丈夫なので、解決案や提案など頂ければ幸いです。

該当のソースコード

↓のmain関数により、とある組み合わせのハイパーパラメータによるGCNの学習とテストをします。
これを1セットとして、最適なハイパーパラメータの組み合わせが見つかるまで、grid-searchによりmain関数を繰り返し実行します。
そのため、main関数が終わる度に、つまりある組み合わせのハイパーパラメータによるGCNの学習とテストが終わる度に、RAM,GPUのメモリをリセットするように動作することが理想です。

from tqdm import tqdm
import hydra
from omegaconf import DictConfig
from hydra import utils

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_scatter import scatter
from torch_geometric.data import RandomNodeSampler
from torch_geometric.nn import GCNConv
from ogb.nodeproppred import PygNodePropPredDataset, Evaluator


# GCNのモデルの定義
class GCN(nn.Module):
    def __init__(self, cfg):
        super(GCN, self).__init__()
        self.dropout = cfg.NN.dropout

        self.convs = nn.ModuleList()
        self.convs.append(GCNConv(cfg.Dataset.n_feat, cfg.NN.n_hid))
        for _ in range(1, cfg.NN.n_layer-1):
            self.convs.append(GCNConv(cfg.NN.n_hid, cfg.NN.n_hid))
        self.convs.append(GCNConv(cfg.NN.n_hid, cfg.Dataset.n_class))

    def forward(self, x, edge_index):
        for l, conv in enumerate(self.convs):
            x = conv(x, edge_index)
            if l < len(self.convs)-1:  # skips relu activate and cfg.dropout iff last layer
                x = F.relu(x)
                x = F.dropout(x, self.dropout, training=self.training)

        return x


def train(epoch, cfg, loader, model, optimizer, device):
    # train
    model.train()
    criteria = torch.nn.BCEWithLogitsLoss()

    num_batches = len(loader)
    for batch_id, data in enumerate(loader): # in [g1, g2, ..., g20]
        data = data.to(device)
        optimizer.zero_grad()
        out = model(data.x, data.edge_index)
        loss = criteria(out[data.train_mask], data.y[data.train_mask])
        loss.backward()
        optimizer.step()


@torch.no_grad()
def test(cfg, loader, model, evaluator, device):
    model.eval()

    ys, preds = [], []
    for data in loader: # only one graph (=g1+g2)
        data = data.to(device)
        out = model(data.x, data.edge_index)
        mask = data['test_mask']
        ys.append(data.y[mask].cpu())
        preds.append(out[mask].cpu())

    test_rocauc = evaluator.eval({
        'y_true': torch.cat(ys, dim=0),
        'y_pred': torch.cat(preds, dim=0),
    })['rocauc']

    return test_rocauc.item()


@hydra.main(config_path='conf', config_name='config')
def main(cfg: DictConfig): # cfgは全てのハイパーパラメータを辞書型として保持する変数
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    root = utils.get_original_cwd() + '/data/ogbn_proteins'
    
    # データセットのローディング
    dataset = PygNodePropPredDataset('ogbn-proteins', root)
    splitted_idx = dataset.get_idx_split()
    data = dataset[0]
    data.node_species = None
    data.y = data.y.to(torch.float)
    
    # データセットの前処理
    row, col = data.edge_index
    data.x = scatter(data.edge_attr, col, 0, dim_size=data.num_nodes, reduce='add')
    cfg.Dataset.n_feat = cfg.Dataset.e_feat

    # 各サンプル(今回の場合頂点)を学習とテストどちらに使うかの印を'mask'としてつける
    for split in ['train', 'valid', 'test']:
        mask = torch.zeros(data.num_nodes, dtype=torch.bool)
        mask[splitted_idx[split]] = True
        data[f'{split}_mask'] = mask

    # ミニバッチをつくる
    train_loader = RandomNodeSampler(data, num_parts=40, shuffle=True,
                                     num_workers=5)
    test_loader = RandomNodeSampler(data, num_parts=5, num_workers=5)

    # 学習とテストの開始
    model = GCN(cfg).to(device)
    optimizer = torch.optim.Adam(params       = model.parameters(), 
                                 lr           = cfg.NN.learning_rate)
    evaluator = Evaluator('ogbn-proteins')

    for epoch in tqdm(range(1, cfg.NN.epochs)):
        train(epoch, cfg, train_loader, model, optimizer, device)
    acc = test(cfg, test_loader, model, evaluator, device)
    print('acc: {:.2f}'.format(acc))

    return acc
    

if __name__ == "__main__":
    main()

↑上記のpythonスクリプトで使われるハイパーパラメータは↓このconfig.yamlファイルで設定しています

defaults:
  - override hydra/sweeper: basic

Dataset: &Dataset
  n_feat: 1
  e_feat: 8
  n_class: 112

NN: &NN
  learning_rate: 0.01
  epochs: 20
  n_hid: 256
  n_layer: 3
  dropout: 0.

↓のようにターミナルで実行すると、hydraというライブラリにより、指定された候補の中から最適なハイパーパラメータの組み合わせを、grid-searchで探索してくれます。

python3 train.py -m 'NN.learning_rate=choice(0.01,0.001)' 'NN.n_hid=choice(64,128,256)' 'NN.n_layer=range(3,6)'

補足情報

上記のソースコードでは、ミニバッチ学習を行っています。また、このメモリがほとんど解放されないまま学習を繰り返しOut of Memoryを引き起こす問題は、ミニバッチ学習の時限定で置きます。バッチ学習時では毎回ちゃんとRAMメモリが解放されました。

試したこと

下記のように、ガーベージコレクターやpytorchのcudaのキャッシュを消したりなどすると、GPUのメモリについては解放されましたがRAMのメモリは依然として解放されませんでした。

dataset = PygNodePropPredDataset('ogbn-proteins', root)
# main関数の最後に下記を実行
del dataset
torch.cuda.empty_cache()
gc.collect()