BigQuery ML Tensorflow 模型 - UDF 内存不足答案

【问题标题】：BigQuery ML Tensorflow model - UDF out of memoryBigQuery ML Tensorflow 模型 - UDF 内存不足
【发布时间】：2021-09-27 16:37:48
【问题描述】：

我正在尝试在 BigQuery 中运行 TensorFlow 模型。该模型是 BERT 的变体，小到足以适应 BigQuery 模型限制 (

我已尝试直接在 BigQuery 控制台中使用以下查询使用模型生成预测：

SELECT
  input_1,
  input_2,
  prediction,
FROM
  ML.PREDICT(MODEL `MY_IMPORTED_MODEL`,
    (
    SELECT
      *,
    FROM
      `MY_DATA_TABLE`
    ))

但是，查询导致以下错误：

查询执行期间资源超出：UDF 内存不足。

我尝试使用以下查询从MY_DATA_TABLE 生成较小样本的预测：

SELECT
  input_1,
  input_2,
  prediction,
FROM
  ML.PREDICT(MODEL `MY_IMPORTED_MODEL`,
    (
    SELECT
      *,
    FROM
      `MY_DATA_TABLE`
    LIMIT 10
    ))

较小的样本效果很好。

我认为OVER 表达式可能会通过强制使用更多插槽来解决问题，因此我生成了以下查询（剧透警告：因内存不足错误而失败）：

SELECT
  input_1,
  input_2,
  prediction,
FROM
  ML.PREDICT(MODEL `MY_IMPORTED_MODEL`,
    (
    SELECT
      *,
      FLOOR(CAST(ROW_NUMBER() OVER (ORDER BY input_1) AS decimal) / 10) AS batch_number,
    FROM
      `MY_DATA_TABLE`
    ))

BigQuery 似乎试图一次向模型提供太多行，从而导致批量大小导致内存不足错误。

由于在调用ML.PREDICT 函数时无法指定BATCH_SIZE 参数，我想知道是否有任何其他方法可以获取不会导致内存不足错误的预测。

有 1.8 亿行要运行预测，所以我想在 BigQuery（而不是 GCP AI 平台）中进行。

有什么想法吗？

【问题讨论】：

标签： tensorflow machine-learning google-bigquery

【解决方案1】：

一种选择是将您的模型拆分为多个模型，并为每个模型调用一个单独的ml.predict。

keras 函数式 api 使这变得非常容易，因为您可以组合不同的模型并在以后使用 get_layer 方法访问它们。

【讨论】：