深度学习项目：如何用Matplotlib快速生成数百万张PNG图像？

问题背景 #

我在深度学习项目中需要为数据集生成250万张224×224像素的合成绘图。目前本机运行100次生成耗时2.7秒，按此速度完成全部任务需约18小时，寻求可行的提速方法。

当前使用的代码：

from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
import matplotlib.pyplot as plt

for i in range(100):
    fig = plt.Figure(frameon=False, facecolor="white", figsize=(4, 4))
    ax = fig.add_subplot(111)
    ax.axis('off')
    ax.plot([1, 2, 3, 4, 5, 6, 7, 8], [2, 4, 6, 8, 8, 6, 4, 3])
    canvas = FigureCanvas(fig)
    canvas.print_figure(str(i), dpi=56)

示例生成图像：

说明：实际生成会使用200行数据，但该因素对生成速度影响极小。

可行提速方案 #

1. 复用Matplotlib核心对象，避免重复初始化 #

每次循环创建fig、ax是主要性能开销点。提前初始化一次画布和轴，循环中仅更新数据：

from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
import matplotlib.pyplot as plt

# 提前完成初始化
fig = plt.Figure(frameon=False, facecolor="white", figsize=(4, 4))
ax = fig.add_subplot(111)
ax.axis('off')
canvas = FigureCanvas(fig)
line, = ax.plot([])  # 初始化空折线

for i in range(100):
    # 更新折线数据
    line.set_data([1, 2, 3, 4, 5, 6, 7, 8], [2, 4, 6, 8, 8, 6, 4, 3])
    # 若数据范围固定可省略以下两行
    ax.relim()
    ax.autoscale_view()
    # 保存图像
    canvas.print_figure(str(i), dpi=56)

该方案可减少30%-50%的单循环耗时。

2. 替换为轻量绘图库 #

Matplotlib偏向交互式绘图，批量生成场景下冗余开销大，可尝试以下替代：

• Pillow + NumPy：直接操作像素数组绘制折线，完全绕过Matplotlib抽象层，速度提升数倍。核心思路：创建224×224的空白NumPy数组，通过数值计算标记折线像素点，再转成Pillow图像保存。
• pycairo：底层绘图库，直接生成像素数据，性能远优于Matplotlib，适合高性能需求场景。
• Plotly Orca：无头渲染器，批量生成速度比Matplotlib快，支持复杂绘图。

3. 多进程并行生成 #

利用多核CPU优势，通过multiprocessing拆分任务：

from multiprocessing import Pool
import matplotlib.pyplot as plt
from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas

def generate_image(i):
    fig = plt.Figure(frameon=False, facecolor="white", figsize=(4, 4))
    ax = fig.add_subplot(111)
    ax.axis('off')
    ax.plot([1, 2, 3, 4, 5, 6, 7, 8], [2, 4, 6, 8, 8, 6, 4, 3])
    canvas = FigureCanvas(fig)
    canvas.print_figure(str(i), dpi=56)
    plt.close(fig)  # 关闭图形避免内存泄漏

if __name__ == '__main__':
    with Pool(processes=4) as pool:  # 进程数设为CPU核心数
        pool.map(generate_image, range(100))

注意：每个进程需独立创建图形对象，且及时关闭避免内存泄漏，总耗时可压缩至单进程的1/N左右（N为进程数）。

4. 优化图像保存流程 #

canvas.print_figure默认设置有冗余开销，可调整为：

import numpy as np
from PIL import Image

# 循环内生成图像后：
canvas.draw()
width, height = fig.get_size_inches() * fig.dpi
image_arr = np.frombuffer(canvas.tostring_rgb(), dtype='uint8').reshape(int(height), int(width), 3)
Image.fromarray(image_arr).save(f"{i}.png")

直接生成内存像素数据再用Pillow保存，比print_figure更高效；同时禁用bbox_inches='tight'等不必要参数，减少计算开销。

5. 硬件加速（可选） #

若有GPU，可尝试：

• TensorFlow绘图函数：直接在GPU上生成图像张量，写入TFRecord格式避免磁盘IO，适合深度学习训练场景。
• CuPy加速：用CuPy替代NumPy处理数值计算，配合Matplotlib可小幅提升速度。

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