Google BigQuery中AUTOML_REGRESSOR预测时TensorFlow Worker内存不足求助

问题重现 #

执行预测时触发以下错误：

google.api_core.exceptions.BadRequest: 400 Resources exceeded during query execution: TENSORFLOW worker out of memory.
This model used too much memory when loaded into RAM to run predictions.
This model may be too complex to run efficiently inside of BigQuery.

即使将预测任务拆分为每批1万行，错误依然存在，说明问题核心并非单批数据量，而是模型本身的内存占用或执行机制导致的。

可行解决方案 #

• 简化模型复杂度

  • 分析模型的特征重要性，移除权重极低的冗余特征，减少模型加载时的内存开销
  • 重新训练模型时，调整AutoML的训练参数：降低max_model_size限制，或限制树模型的深度/数量（若模型基于树结构）
  • 尝试切换至更轻量化的模型（如线性回归），在精度可接受的前提下优先选择低内存占用的模型类型

• 调整预测执行策略

  • 导出模型至Vertex AI或本地TensorFlow环境，利用分布式计算框架处理大批次预测，避开BigQuery worker的内存限制
  • 在BigQuery内执行时，申请更大的查询槽分配：通过SET query_slot_limit = [更大数值]提升单worker的内存配额
  • 优化分批逻辑：不要仅按行数均分，结合数据特征的稀疏性分组，避免某批数据包含大量高维度稀疏特征导致内存骤增

• 优化预测数据集

  • 对高基数类别特征做哈希分桶或降维处理，减少特征维度带来的内存占用
  • 标准化所有数值特征，避免因数值范围差异引发的模型计算内存波动
  • 过滤预测数据中的无效记录（如空值占比过高的行），减少不必要的计算负载

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