BigQuery ML Tensorflow model - memory 中的 UDF

Question

我正在尝试在 BigQuery 中运行 Tensorflow model。 model 是 BERT 的变体，它足够小以符合 BigQuery model 的限制 (<250MB)。

我尝试直接从 BigQuery 控制台使用以下查询生成 model 的预测：

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

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

查询执行期间超出资源：memory 中的 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 试图一次向 model 提供太多行，导致一批大小导致内存不足错误。

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

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

有任何想法吗？

Answer 1

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

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