使用shap.KernelExplainer计算Keras版BERT分类器SHAP值时出现OOM张量分配错误

ResourceExhaustedError: OOM when allocating tensor with shape[23020,128,768] and type float on /job:localhost/replica:0/task:0/device:GPU:0 by allocator GPU_0_bfc [Op:Tile]

• SHAP的KernelExplainer属于模型无关解释工具，计算过程中会为每个输入样本生成大量扰动样本，即使仅传入3条测试样本，最终批量推理的总样本量也达到了2万+（对应报错张量的第一维大小23020），叠加BERT单样本推理生成的[128,768]大小隐层张量，显存占用直接超出GPU上限。
• 模型调用参数不匹配：你自定义的Classifier类的call方法需要同时接收input_ids和attention_mask两个参数，但你仅将input_ids传入了KernelExplainer，扰动样本无对应的attention_mask，导致模型产生额外无效计算，进一步占用显存。
• TensorFlow默认占用全部GPU显存，且没有显式开启推理模式，也会加剧显存占用压力。

按可落地优先级排序：

1. 更换更适配NLP场景的SHAP解释器
   放弃KernelExplainer，改用PartitionExplainer，该解释器专门针对序列数据设计，扰动样本生成量远低于KernelExplainer，显存占用可降低70%以上。如果必须使用KernelExplainer，调用shap_values时添加nsamples参数限制扰动样本量，示例：

kernel_shap_values = kernel_explainer.shap_values(X_test, nsamples=50)

默认nsamples为几百，可根据显存情况逐步调低。
2. 包装模型统一输入
将模型包装为仅接收input_ids单个参数的推理函数，内部自动生成对应attention_mask，避免参数不匹配导致的无效显存占用：

def predict_wrapper(input_ids):
    attention_mask = tf.cast(input_ids != tokenizer.pad_token_id, tf.int32)
    return model(input_ids=input_ids, attention_mask=attention_mask, training=False)
# 用包装后的函数初始化explainer
kernel_explainer = shap.KernelExplainer(predict_wrapper, X_train)

3. 配置TensorFlow显存动态分配
   在代码最开头添加显存配置，避免TensorFlow预先占用全部GPU显存：

import tensorflow as tf
gpus = tf.config.list_physical_devices('GPU')
if gpus:
    for gpu in gpus:
        tf.config.experimental.set_memory_growth(gpu, True)

4. 进一步降低显存占用的可选操作

• 继续缩减背景样本量，从当前的10条降到3-5条，可直接降低KernelExplainer的计算量
• 调低max_len参数，根据实际文本平均长度设置为32或64，单样本张量大小会线性下降
• 若以上操作都无法解决，可设置环境变量CUDA_VISIBLE_DEVICES=-1强制使用CPU计算，虽然速度较慢但不会触发显存不足报错

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