BigQuery ML Tensorflow model - memory 中的 UDF

Question

I'm trying to run a Tensorflow model in BigQuery. The model is a variant of BERT that is small enough to fit within BigQuery model limitations (<250MB).

I've tried to generate the predictions with the model using the following query directly from the BigQuery console:

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

However, the query resulted in the following error:

Resources exceeded during query execution: UDF out of memory.

I have attempted to generate the predictions on a smaller sample from MY_DATA.TABLE with the following query:

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

The smaller sample works perfectly fine.

I thought that maybe an OVER expression would fix the issue by forcing the usage of more slots, so I produced the following query (spoiler alert: failed with the out-of-memory error):

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`
    ))

It seems that BigQuery attempts to feed too many rows at once to the model, resulting in a batch of the size that results in the out-of-memory error.

Since I cannot specify the BATCH_SIZE parameter while calling the ML.PREDICT function, I'd like to know if there's any other way of obtaining the predictions that wouldn't result in the out-of-memory errors.

There are 180 million rows to run the prediction for, so I'd like to do it from within the BigQuery (not the GCP AI Platform).

Any ideas?

Answer 1

One option is to split your model into multiple models and call a separate ml.predict for each of them.

The keras functional api makes this quite easy as you can compose different models and access them later with the get_layer method.