循环生成并保存Matplotlib图后如何清理内存避免内存错误？

问题背景 #

我有8000个时间步的HDF5格式模拟数据，将其加载为HistoryHdf5类对象后，通过循环遍历生成Matplotlib图表并保存至磁盘。但执行约100次循环后代码因内存错误崩溃，尝试过相关解决方案但无效，现寻求循环后清理内存与图形的有效方法，原代码如下：

# 加载HDF5文件
history = HistoryHdf5.from_archive("CW_0.5_1.hf5")

# 遍历所有时间步
for times in np.arange(0, 800, 0.1):

    sheet = all_quartet_composition_edge(history, times)

    identify_rosettes(sheet)

    draw_specs = tyssue.config.draw.sheet_spec()

    color = sheet.edge_df["composition"]
    sheet.face_df["color"] = 0
    sheet.face_df["chirality"] = "CW"
    sheet.face_df.loc[sheet.face_df["torque_coef"] == 0.2, "color"] = 1
    sheet.face_df.loc[sheet.face_df["torque_coef"] == 0.2, "chirality"] = "CCW"
    cmap = mpl.colors.ListedColormap(["black", 'red', 'blue', 'purple', 'orange', 'yellow'])
    cmap2 = mpl.colors.ListedColormap(["white", "black"])
    color_map = cmap(color)
    color_map2 = cmap2(sheet.face_df["color"])
    color_map3 = cmap2(sheet.vert_df["rosette"])
    draw_specs['face']['color'] = "white"
    draw_specs['face']['visible'] = True
    draw_specs['face']['alpha'] = 0.2
    draw_specs['vert']['visible'] = True
    draw_specs['vert']['color'] = color_map3
    draw_specs['vert']['alpha'] = sheet.vert_df["rosette"]
    draw_specs['edge']['color'] = color_map
    draw_specs['edge']['head_width'] = 0.01
    draw_specs['edge']['width'] = 2
    fig, (ax) = plt.subplots(1, figsize = (12, 12), dpi = 300)
    fig, ax = sheet_view(sheet, ["x", "y"], ax, **draw_specs)

    black_patch = mpatches.Patch(color='black', label='Boundary/Rosette')
    red_patch = mpatches.Patch(color='red', label='Homogenous')
    blue_patch = mpatches.Patch(color='blue', label='3:1 Heterogenous')
    purple_patch = mpatches.Patch(color='purple', label='3:1 Homogenous')
    orange_patch = mpatches.Patch(color='orange', label='2:2 Heteroogenous')
    yellow_patch = mpatches.Patch(color='yellow', label='2:2 Homogenous')

    plt.legend(handles=[black_patch, red_patch, blue_patch, purple_patch, orange_patch, yellow_patch])

    for face, data in sheet.face_df.iterrows():
        ax.text(data.x-0.3, data.y, sheet.face_df["chirality"][face], fontsize=4, color="b")
    
    fig.savefig(f"test_{times}.png")

    plt.cla()
    plt.clf()
    plt.close(fig)

有效解决方案 #

针对内存泄漏的核心问题，以下是经过验证的优化措施：

• 彻底清理Matplotlib资源：放弃全局plt函数的清理方式，改用面向对象的资源管理，循环末尾显式删除图表、轴对象，避免残留引用
• 回收模拟数据对象：每次循环生成的sheet对象包含大量DataFrame数据，循环末尾必须显式删除并触发垃圾回收
• 复用静态资源：将颜色映射表（cmap）、图例补丁等不会随时间步变化的对象移到循环外，避免重复创建浪费内存
• 优化DataFrame操作：合并重复的DataFrame赋值逻辑，减少临时内存占用
• 强制垃圾回收：循环末尾调用gc.collect()，主动回收Python解释器未自动释放的内存

修改后的代码 #

import gc
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
import tyssue

# 加载HDF5文件
history = HistoryHdf5.from_archive("CW_0.5_1.hf5")

# 提前创建静态资源，避免循环内重复创建
cmap = mpl.colors.ListedColormap(["black", 'red', 'blue', 'purple', 'orange', 'yellow'])
cmap2 = mpl.colors.ListedColormap(["white", "black"])
# 预创建图例补丁
legend_patches = [
    mpatches.Patch(color='black', label='Boundary/Rosette'),
    mpatches.Patch(color='red', label='Homogenous'),
    mpatches.Patch(color='blue', label='3:1 Heterogenous'),
    mpatches.Patch(color='purple', label='3:1 Homogenous'),
    mpatches.Patch(color='orange', label='2:2 Heteroogenous'),
    mpatches.Patch(color='yellow', label='2:2 Homogenous')
]

# 遍历所有时间步
for times in np.arange(0, 800, 0.1):
    sheet = all_quartet_composition_edge(history, times)
    identify_rosettes(sheet)

    draw_specs = tyssue.config.draw.sheet_spec()

    color = sheet.edge_df["composition"]
    # 合并DataFrame赋值操作，减少内存波动
    sheet.face_df["color"] = 0
    sheet.face_df["chirality"] = "CW"
    ccw_mask = sheet.face_df["torque_coef"] == 0.2
    sheet.face_df.loc[ccw_mask, ["color", "chirality"]] = [1, "CCW"]

    color_map = cmap(color)
    color_map2 = cmap2(sheet.face_df["color"])
    color_map3 = cmap2(sheet.vert_df["rosette"])
    
    draw_specs['face']['color'] = "white"
    draw_specs['face']['visible'] = True
    draw_specs['face']['alpha'] = 0.2
    draw_specs['vert']['visible'] = True
    draw_specs['vert']['color'] = color_map3
    draw_specs['vert']['alpha'] = sheet.vert_df["rosette"]
    draw_specs['edge']['color'] = color_map
    draw_specs['edge']['head_width'] = 0.01
    draw_specs['edge']['width'] = 2

    # 面向对象创建图表，避免全局plt状态污染
    fig, ax = plt.subplots(1, figsize=(12, 12), dpi=300)
    fig, ax = sheet_view(sheet, ["x", "y"], ax, **draw_specs)

    ax.legend(handles=legend_patches)

    # 优化文本绘制，直接用data中的chirality列
    for face, data in sheet.face_df.iterrows():
        ax.text(data.x - 0.3, data.y, data.chirality, fontsize=4, color="b")
    
    fig.savefig(f"test_{times}.png")

    # 彻底清理Matplotlib资源
    plt.close(fig)
    # 显式删除对象，解除引用
    del fig, ax, color_map, color_map2, color_map3, draw_specs, sheet, ccw_mask, color
    # 强制垃圾回收
    gc.collect()

内容的提问来源于stack exchange，提问作者drsnif