tensorflow 2.x model.predictのメモリーリーク

はじめに

tensorflow2.xでKerasのModel.predictを繰り返し実行したら、メモリ使用量がもりもり増えていったので、メモ。
Inputサイズが大きいというのもありますが、128Gメモリ、バッチサイズ4000、4回のpreditほどで124G張り付き。。。
ソースまで深追いはしていませんが、現状メモリーリーク(memory leak）があるようです。

ubuntu==18.04
python==3.8.5
tensorflow==2.4.1
nvidia-driver==465.19.01
cuda==11.0

対策：うまく行かなかったパターン

マニュアルにメモリをクリアしてみる。
疑似コードとしては以下のように修正

for x in dataset:
    y_pred = mode.predct(x)

    # 明示的に削除とGC
    del y_pred
    y_pred = None
    gc.collect()

結果としては、多少（数ループ）改善したものの8000回のループ回したい要件としてはNG

対策：うまく行ったパターン

参考にしたのは、↓にリンクされていたgithubのissueをよく読んでみました。
するとpredict_on_batchでは問題なく動作するとのこと

疑似コードとしては以下のように修正

for x in dataset:
    y_pred = mode.predct_on_batch(x)

20〜30Gぐらいで均衡を保つようになりました