Matplotlibのsavefig()でメモリリークする問題を深堀りした

はじめに

Matplotlibでたくさんのグラフを画像として保存するプログラムを書いたところ、メモリリークで落ちるという問題にぶち当たりました。調べてみると何人かの方がこれについて解決策を提示して下さっています。その中ではこれが決定版でしょうか。

ただ、自分の環境ではちょっと設定が違ったせいで上のやり方ではうまく行かず (後述しますが、手法が悪かったわけではなく私のコードのバグでした) そこからさらに深堀したおかげでいろいろわかったことがあります。以下整理してみようと思います。

先に結論

• Matplotlibのsavefig()におけるメモリリーク問題の根本はinlineバックエンドにあり
  • savefig()したいだけならAggを使え
  • バックエンドの選定次第では処理速度にも大きく影響
• グラフの初期化を一度だけやって毎回cla()する方法はオールマイティ
  • ただし複数のAxesを扱う際には注意が必要

環境

以下の環境で確認しました。

• OS: Debian 11.11 (VSCodeのDevContainer環境でmcr.microsoft.com/devcontainers/python:1-3.12-bullseyeを使用)
• Python: 3.12.11
• matplotlib: 3.10.6
• jupyter: 4.4.7

再現実験

@uz29さんの上記の記事のコードを踏襲しつつ、少し修正しました。

import psutil
import shutil
import time
from pathlib import Path

import numpy as np
import matplotlib.pyplot as plt

log_file = Path("log.csv")
fw = log_file.open("w")
fw.write("i,time_delta[s],memory[MB]\n")
outdir = Path("./output")
if outdir.exists():
    shutil.rmtree(outdir)
outdir.mkdir()

max_iter = 1000
memory_limit = 1e9

memory_start = psutil.virtual_memory().used
time_start = time.time()
org_time_start = time_start

for i in range(max_iter):
    # グラフ作成
    size = 100
    y = np.random.randn(size)
    fig, ax = plt.subplots(figsize=(8, 8))
    ax.plot(y)
    
    # 画像出力
    outfile = (outdir / f"{i:05d}").with_suffix(".png")
    plt.savefig(outfile)
    
    # グラフデータを解放する
    # plt.cla()
    # plt.clf()
    plt.close()
    
    # 開始時からのメモリ増加量と1ループの時間を計測
    memory_delta = psutil.virtual_memory().used - memory_start
    time_end = time.time()
    time_delta = time_end - time_start
    time_start = time_end
    out_str =f"{i + 1},{time_delta},{memory_delta / 1e6}"
    print(out_str) 
    fw.write(out_str + "\n")
    if memory_delta > memory_limit:
        print("Memory consumption exceeds the limit.")
        break

fw.close()
total_time = time_end - org_time_start
total_iter = i + 1
avg_time = total_time / total_iter
print(f"average time delta: {avg_time}")

pathlibの使用やf-stringの多用は趣味の問題かも。他に以下の点を修正しています。

• メモリ消費の上限を設定し、それを超えたら終了
• 処理時間の計測にはtime.time()を使用
• グラフは簡略化
• 画像出力先は、条件が同じになるように毎回ディレクトリ削除
• グラフデータ解放時の丁寧バージョンはclf() → cla() → close()ではなくcla() → clf() → close()
  • clf()でfigureをクリアした時点でAxesもクリアされるので、順を追うとすればAxesをクリアしてからfigureをクリアすべきかと

close()のみ版は上記のまま、cla() → clf() → close()については# グラフデータを解放するのあとの2行のコメントアウトを戻して実行しました。さらに、@uz29さんの提案する手法である「最初にグラフを一度だけ初期化して、ループの中ではcla()だけを毎回行う」については以下のようにしています。

# グラフ領域作成
fig, ax = plt.subplots(figsize=(8, 8))

for i in range(max_iter):
    # グラフ作成
    size = 100
    y = np.random.randn(size)
    ax.plot(y)
    
    # 画像出力
    outfile = (outdir / f"{i:05d}").with_suffix(".png")
    plt.savefig(outfile)
    
    # Axesをクリア
    plt.cla()
    
    # 開始時からのメモリ増加量と1ループの時間を計測
    ...

これらをJupyter Notebook上で実行した際のメモリ消費量のグラフは下図のようになりました。

ラベルはお分かりとは思いますが

• close only: ループ内でclose()のみを毎回実施
• cla, clf and close: ループ内でcla() → clf() → close() を毎回実施
• init once, cla each: 最初にグラフを初期化、ループ内でcla()を毎回実施

です。

cla, clf and closeの傾斜がオリジナルの記事より急ですが、傾向としては概ね同様として良いのではないかと思います。

バックエンドの変更

実はいろいろ調べている過程で、GoogleのAI Overviewに「バックエンドがメモリ消費に影響する」といった旨の示唆があるのを見かけました。今回改めて"matplotlib savefig メモリリーク"で検索した際のAIモードの回答から引用します。

対処法3：バックエンドを切り替える
インタラクティブな描画が不要な場合は、Matplotlibのバックエンドを非インタラクティブな Agg に切り替えることで、メモリ使用量が大幅に改善されることがあります。これは、サーバ上での実行や、大量の画像をバッチで生成する場合に特に有効です。

というわけで、バックエンドにAggを指定して同じことをやってみます。バックエンドやAggの詳細に関してはMatplotlibのドキュメントのBackendsを読んでみてください。

ちなみに上のコードではバックエンドを何も指定していませんが、現行のJupyterではデフォルトでinline (module://matplotlib_inline.backend_inline) が選ばれるようです。Aggへの変更は下記のようにすればOK。

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

これ以外は先ほどと同じコードで3種類の方法を試してみます。結果はこちら。

close onlyもcla, clf and closeも、1000回繰り返しても100MB以下に収まっています。init once, cla eachに比べればメモリ消費量自体はだいぶ多いですが、概ね横ばいでこれ以上増えそうな気配はありません。つまり、猛烈な勢いでメモリを食いつぶす原因はinlineバックエンドであり、かわりにAggを指定しておけばclose onlyで大丈夫、と言って差し支えないと思います。

close onlyとcla, clf and closeに差がないことは、本来の意味を考えると納得感があります。API Referenceのmatplotlib.pyplot.closeには

Close a figure window, and unregister it from pyplot.

と書かれており、またオブジェクトの親子関係を考えるとfigureがクローズされた時点でAxesは破棄されると考えるのが自然です。ではinlineではなぜclose()しても破棄されないのか。Jupyter Notebook上で文字通りinline表示を行うためのバックエンドなので、グラフィカルなグラフを描画したら処理の如何に関わらず消えないことが期待され、形の上ではpypoltの手を離れても個々のfigureは一定のメモリを占有するのではないかと思います。

Aggにおけるclose onlyとinit once, cla eachの差ですが、close onlyの場合は毎回figureが生成されており、その破棄はガベージコレクションに回されるので、ある程度メモリは消費されるもののそれ以上は増えない、という状況だと考えられます。一方init once, cla eachではfigureは一つであり、それが毎回クリアされて使いまわされるので、メモリ増加量ははるかに少ないということでしょう。グラフのレイアウトは最初に決めたものから変えられませんが、バックエンドの問題に左右されない点ではオールマイティと言えます。

Jupyterではなく普通のpythonスクリプトの場合

同じコードをエクスポートして通常のpythonスクリプトとして保存し、シェルで実行してみます。このとき、私の環境ではバックエンドにtkAggが使用されていました (matplotlib.get_backend()で確認できます)。close onlyの結果を比較したものがこちら。inlineはJupyter上での実行、AggとtkAggはシェルからの実行です。

tkAggのメモリ消費はAggよりは多いですが、inlineのような爆発は起きません。tkAggの場合はinlineとは違い、本来のグラフィカル表示においてはclose()で画面が破棄されますので、メモリは溜まらないのでしょう。

ただ、実は別の問題があります。画像1枚あたりの平均処理時間は以下の通りでした。

Backend	time delta [s]
inline	0.12
Agg	0.19
tkAgg	1.35

tkAggが突出して遅いです。メモリよりは処理時間の観点で、通常のpythonスクリプトでもAggを指定すべきであることがわかります。もっとも、tkAggでもinit once, cla eachの手順を踏むと処理時間は0.1秒台、メモリ消費量もJupyterでのケースと同程度ではあります。

cla()に関する注意事項

冒頭述べた「自分の環境ではちょっと設定が違ったせいで上のやり方ではうまく行かず」の話です。描画したいグラフが複数のサブプロットを持つケースで注意が必要です。

例えばこんなケース。

# グラフ領域作成
num_plots = 2
fig, axs = plt.subplots(nrows=num_plots, ncols=1, figsize=(8, 8))

for i in range(max_iter):
    size = 100
    for ax in axs:
        y = np.random.rand(size)
        ax.plot(y)

    #画像出力
    outfile = (outdir / f"{i:05d}").with_suffix(".png")
    plt.savefig(outfile)
    
    plt.cla()

これ、実行するとどんどん遅くなります。最初理由がわからなかったのですが、描かれた図を見て一目瞭然。10回目でこうなります。

上のグラフが全くクリアされずどんどん重ね描きされているのです。API ReferenceでMatplotlib.pyplot.claを見ると、

Clear the current Axes.

と書かれていて、カレントのAxesだけがクリアされることがわかります。私の実際のケースでは一つ目のグラフがimshow()を使った画像表示だったため、最後に重ねられた画像だけが有効で一見正常に描画されているように見えて、なかなか気付けませんでした。対策としては、カレントのAxesを順次切り替えてからcla()するしかないようです。上のコードの最後の行を下記のようにすればOKです。sca()はset current Axesですね。

    for ax in axs:
        plt.sca(ax)
        plt.cla()

まとめ

というわけで冒頭述べた結論の繰り返しです。

• Matplotlibのsavefig()におけるメモリリーク問題の根本はinlineバックエンドにあり
  • savefig()したいだけならAggを使え
  • バックエンドの選定次第では処理速度にも大きく影響
• グラフの初期化を一度だけやって毎回cla()する方法 (init once, cla each) はオールマイティ
  • ただし複数のAxesを扱う際には注意が必要

実際のところ、グラフィカル表示が必要ないならAggを使うことは習慣化しておいたほうが良いと思います。その上で、ラフに考えるならclose()さえ忘れなければ大きな問題には至らないでしょう。一方で、何等かの事情でAggが使えないときや原因不明のメモリリークが発生した時 (importしている他のモジュールがバックエンドを上書きしているようなケースはあるかもしれません) ではinit once, cla eachが解決してくれる可能性が高いです。その際は、クリアすべきAxesが複数ある時のクリア漏れに注意する必要があります。

以上、本件についてわかったことをまとめてみました。何かのお役に立てれば嬉しいです。誤記、誤解の指摘や質問、コメント等は大歓迎です。よろしくお願いします。