Python：如何提升批量shap依赖关系图导出为单个PDF的速度

核心耗时原因分析 #

你的代码里存在大量无必要的交互场景开销，加上SHAP默认的额外计算逻辑，导致整体耗时被拉长，以下是可落地的优化方案：

可直接落地的优化点 #

• 删除无用暂停逻辑
  plt.pause(0.1)是为交互场景下刷新画布设计的，批量生成图片完全不需要该逻辑，直接删除即可，2000张图可直接减少200秒无意义等待。
• 切换非交互渲染后端
  matplotlib默认的交互后端会额外计算界面渲染相关的内容，批量生成时切换为无UI的Agg后端，可降低30%左右的渲染开销，只需在代码最开头加入：

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

• 禁用SHAP自动交互项计算
  shap.dependence_plot默认interaction_index='auto'，会自动为每个特征计算最优交互特征并染色，该计算逻辑占单张图生成耗时的近一半，如果你不需要展示交互项，直接设置interaction_index=None即可。
• 复用画布减少对象创建开销
  避免每次生成新图都创建、销毁Figure对象，提前创建固定大小的画布，每次画图前清空内容，可减少大量对象初始化开销。
• 去掉自适应边距计算
  pdf.savefig的bbox_inches = 'tight'会额外计算每张图的边距，如果你可以接受固定边距，提前设置画布边距后去掉该参数，可进一步提升保存速度。
• 用普通循环替代pandas apply
  pandas的apply有额外的调度开销，直接遍历列名执行函数即可。

优化后单进程代码示例 #

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.backends.backend_pdf import PdfPages
import shap
import pandas as pd

def plot_shap_dependence(col_txt, shap_values, x_test, fig, ax):
    # 清空画布
    ax.clear()
    shap.dependence_plot(
        col_txt, 
        shap_values,
        x_test,
        show=False,
        interaction_index=None, # 不需要交互项就保留，需要的话删掉该行
        ax=ax # 复用已有的坐标轴
    )
    return fig

# 初始化配置
show = False
shapValues_test = shapTreeExplainer.shap_values(X_test)
# 提前创建固定大小的画布和坐标轴
fig, ax = plt.subplots(figsize=(8, 5))

with PdfPages('output.pdf') as pdf:
    for col in X_test.columns:
        current_fig = plot_shap_dependence(col, shapValues_test[1], X_test, fig, ax)
        pdf.savefig(current_fig)
# 最后只关一次画布即可
plt.close(fig)

进阶多核优化方案 #

如果你的CPU是多核的，可将特征列表拆分为多份，用多进程并行生成独立的PDF片段，最后合并为一个PDF：

1. 用multiprocessing.Pool启动多进程，每个进程处理一部分特征，生成单独的子PDF
2. 用PyPDF2或者pikepdf把所有子PDF合并为一个完整文件，8核CPU可将总耗时压缩到原来的1/5~1/7。

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