[Pytorch / ResNet画像分類 / テストデータ予測] RAMメモリがfold毎に増加する・メモリが解放されない

Question

前提

Pytorchで画像分類をしています。（ResNet200Dモデル）

交差検証で事前に訓練したモデルが5fold分あり、
それらのモデルを使ってテストデータの予測を5つ分行い、
アンサンブルをしようとしています。

発生している問題

下記ソースコード中の、inference()を実行すると、
各fold毎に、RAMメモリが増加していき、
最終的にメモリが解放されないまま終了して困っています。

具体的には、RAMメモリの使用状況が、

(実行前) 0.2 GB
(1ループ目）2.5 GB
(2ループ目) 3.3 GB
・・・
(5ループ目) 6.0 GB
(実行後) 5.5 GB

のように遷移しています。

GPUメモリについても、2GBほど増えたまま解放されません。

・fold毎にRAMメモリ使用量が増加するのはなぜか。また、その解決策。
・実行後に、RAMメモリ・GPUメモリが解放されないのはなぜか。また、その解決策。

以上、ご教示いただければ幸いです。

該当のソースコード

device = torch.device('cuda')
df_test = pd.read_csv(CSV_PATH)

def inference_one_loop(model, loader, device):

    model.to(device)
    model.eval()
    pred_list = []
    with torch.no_grad():
        for x in tqdm(loader):
            y = model(x.to(device))
            pred_list.append(y.sigmoid().detach().cpu().numpy())
                    
    pred_arr = np.concatenate(pred_list)

    # del pred_list
    # torch.cuda.empty_cache()
    # gc.collect()

    return pred_arr


def inference():
    
    n_fold = 5
    n_class = 11
    batch_size = 32
    image_size = 512

    Model = ResNet200D
    model_name = 'resnet200d'
    model_paths = ['./fold0.pth','./fold1.pth','./fold2.pth','./fold3.pth','./fold4.pth','./fold5.pth']

    test_dataset = TestDataset(df_test, transform=get_transforms(image_size=image_size))
    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=4 , pin_memory=True, drop_last=False)

    ### 各foldごとにメモリが増加する ###

    test_preds_arr = np.zeros((n_fold, len(df_test), n_class))    
    for fold in range(n_fold):
        logger.info(f"[fold {fold}]")

        model = Model(model_name=model_name)
        model.load_state_dict(torch.load(model_paths[fold], map_location=device))

        test_pred = inference_one_loop(model, test_loader, device)
        test_preds_arr[fold] = test_pred

        # del model
        # torch.cuda.empty_cache()
        # gc.collect()

    return test_preds_arr.mean(axis=0)


### 関数を抜けてもメモリが解放されない。
preds = inference()

試したこと

ソースコード中でコメントアウトしているように、
del による変数の参照の削除や、
torch.cuda.empty_cache()を実施しましたが、
結果は変わりませんでした。

また、下記の記事を参考にしましたが、
メモリを多く占有している変数は見つかりませんでした。
https://qiita.com/AnchorBlues/items/883790e43417640140aa

inference()の実行後は、
メソッド内のローカル変数を参照することは当然ながらできないようです。

補足情報（FW/ツールのバージョンなど）

環境: Jupyter Notebook (Kaggle上のnotebook環境）

ライブラリ:
・torch 1.7.0
・timm (https://github.com/rwightman/pytorch-image-models)

モデルやDatasetの定義:

class ResNet200D(nn.Module):
    def __init__(self, model_name):
        super().__init__()
        self.model = timm.create_model(model_name, pretrained=False)
        n_features = self.model.fc.in_features
        self.model.global_pool = nn.Identity()
        self.model.fc = nn.Identity()
        self.pooling = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Linear(n_features, 11)

    def forward(self, x):
        bs = x.size(0)
        features = self.model(x)
        pooled_features = self.pooling(features).view(bs, -1)
        output = self.fc(pooled_features)
        return output

def get_transforms(image_size=640):
    return Compose([
        Resize(image_size, image_size),
        Normalize(),
        ToTensorV2(),
    ])

class TestDataset(Dataset):
    def __init__(self, df, transform=None):
        self.df = df
        self.file_names = df['StudyInstanceUID'].values
        self.transform = transform
        
    def __len__(self):
        return len(self.df)

    def __getitem__(self, idx):
        file_name = self.file_names[idx]
        file_path = f'{TEST_DIR}/{file_name}.jpg'
        image = cv2.imread(file_path)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        if self.transform:
            augmented = self.transform(image=image)
            image = augmented['image']
        return image

Answer 1

tqdmへのイテレータの渡し方が問題でした。

tqdm(loader)

ではなく、

tqdm(enumerate(loader), total=len(loader))

とするのが正しいです。

こちらの質問が参考になりました。
https://github.com/tqdm/tqdm/issues/746

また、問題の発見には、
sys.getrefcount()で、objectの参照数を調べることが役に立ちました。

inference()関数を抜ける直前に、すべてのローカル変数について参照数を調べたところ、
test_loader変数のみ参照が残っていることがわかりました。