はじめに

BigQuery ML はインポートした TensorFlow モデルでの予測ができます。
BigQuery ML を使って TensorFlow モデルを管理すると、データの転送を省略したり、モデルを BigQuery と Cloud Storage に一任できます。
この記事では、BigQuery ML にインポート可能な TensorFlow モデルを作っていきます。

なお、この記事は @K_Urushi のご協力で作成しています。

モデルの作り方

TensorFlow モデルをインポートする CREATE MODEL ステートメントにあるように、BigQuery ML で使える TensorFlow モデルは SavedModel として保存されている必要があります。

SavedModel を作る

シンプルな SavedModel を作る

tf.saved_model.save の例から始めましょう。
例では、tf.TensorSpec は shape=None と定義されていますが、BigQuery ML から使う場合は必須のようですので、
shape=1 とします。関数を変える場合には、適切に tf.TensorSpec を書き換えてください。

import tensorflow as tf


class Adder(tf.Module):

  @tf.function(input_signature=[tf.TensorSpec(shape=1, dtype=tf.float32)])
  def add(self, x):
    return x + x + 1.

to_export = Adder()
tf.saved_model.save(to_export, '/tmp/adder')
# 認証しておけば Google Storage に直接転送できます
# tf.saved_model.save(to_export, 'gs://tmp/adder')

作った SavedModel を BigQuery にインポートする

TensorFlow モデルのインポートを参考にモデルをインポートします。
Cloud Storage にある SavedModel を参照できるので、予め転送しておきましょう。
クエリ 1 つでインポートできるのでとてもお手軽です。

CREATE OR REPLACE MODEL
  example_dataset.imported_tf_model OPTIONS (MODEL_TYPE='TENSORFLOW',
    MODEL_PATH='gs://tmp/adder/*')

インポートしたモデルを使う

インポートした TensorFlow モデルでの予測を参考にモデルで予測します。

SELECT
  *
FROM
  ML.PREDICT(MODEL example_dataset.imported_tf_model,
    (
    SELECT
      *
    FROM
      UNNEST(GENERATE_ARRAY(1,10))x

実行結果

行	output_0	x
0	3.0	1
1	5.0	2
2	7.0	3
3	9.0	4
4	11.0	5
5	13.0	6
6	15.0	7
7	17.0	8
8	19.0	9
9	21.0	10

無事に実行できました。

tf.function の特徴である AutoGraph によって、Python コードが TensorFlow グラフに変換可能です。JavaScript UDF の代わりに使えるかもしれません。

サポートされている型

サポートされている入力にありますが、再掲します。

TensorFlow 型	BigQuery ML type
tf.int8, tf.int16, tf.int32, tf.int64, tf.uint8, tf.uint16, tf.uint32, tf.uint64	INT64
tf.float16, tf.float32, tf.float64, tf.bfloat16	FLOAT64
tf.bool	BOOL
tf.string	STRING

2020 年 2 月 12 日現在、対応している入出力型は限定的なため、BigQuery のデータ型とモデル作成時の型の自由度の差異に注意しましょう。

tf.estimator を使って SavedModel を作る

続いて、予測に使用する SavedModel のエクスポートを参考に、機械学習モデルの SavedModel を作っていきます。

ここでは、BigQuery ML のリリースが待たれる Boosted Trees として、Boosted trees using Estimators を見ながら、tf.estimator.BoostedTreesClassifier を作ってみましょう。

データのロード

import numpy as np
import pandas as pd
import tensorflow as tf


dftrain = pd.read_csv('https://storage.googleapis.com/tf-datasets/titanic/train.csv')
dfeval = pd.read_csv('https://storage.googleapis.com/tf-datasets/titanic/eval.csv')
y_train = dftrain.pop('survived')
y_eval = dfeval.pop('survived')

入力値の作成

fc = tf.feature_column
# カテゴリ値に対応すると煩雑になったため省略
# CATEGORICAL_COLUMNS = ['sex', 'n_siblings_spouses', 'parch', 'class', 'deck',
#                        'embark_town', 'alone']
NUMERIC_COLUMNS = ['age', 'fare']

def one_hot_cat_column(feature_name, vocab):
  return tf.feature_column.indicator_column(
      tf.feature_column.categorical_column_with_vocabulary_list(feature_name,
                                                 vocab))
feature_columns = []
for feature_name in CATEGORICAL_COLUMNS:
  # Need to one-hot encode categorical features.
  vocabulary = dftrain[feature_name].unique()
  feature_columns.append(one_hot_cat_column(feature_name, vocabulary))

for feature_name in NUMERIC_COLUMNS:
  feature_columns.append(tf.feature_column.numeric_column(feature_name,
                                           dtype=tf.float32))


NUM_EXAMPLES = len(y_train)

def make_input_fn(X, y, n_epochs=None, shuffle=True):
  def input_fn():
    dataset = tf.data.Dataset.from_tensor_slices((dict(X), y))
    if shuffle:
      dataset = dataset.shuffle(NUM_EXAMPLES)
    # For training, cycle thru dataset as many times as need (n_epochs=None).
    dataset = dataset.repeat(n_epochs)
    # In memory training doesn't use batching.
    dataset = dataset.batch(NUM_EXAMPLES)
    return dataset
  return input_fn

# Training and evaluation input functions.
train_input_fn = make_input_fn(dftrain, y_train)
eval_input_fn = make_input_fn(dfeval, y_eval, shuffle=False, n_epochs=1)

tf.estimator の作成

est = tf.estimator.BoostedTreesClassifier(feature_columns,
                                          n_batches_per_layer=1)

est.train(train_input_fn)

# Eval.
# result = est.evaluate(eval_input_fn)
# print(pd.Series(result))

SavedModel の作成

json_serving_input_fn を作って、export すると BigQuery ML から理想的な形で呼び出すことができます。

TensorFlow 1.x

tf.placeholder が使えるのでこのコードが動作します。

def json_serving_input_fn():
    """Build the serving inputs."""
    inputs = {}
    for feat in feature_columns:
        print(feat)
        inputs[feat.name] = tf.placeholder(shape=[None], dtype=feat.dtype)

    return tf.estimator.export.ServingInputReceiver(inputs, inputs)

path = est.export_saved_model('gs://tmp/btc',
                      json_serving_input_fn)

TensorFlow 2.x

tf.placeholder の代わりに、tf.compat.v1.placeholder を使えます（@hidekiy 氏より）。


def json_serving_input_fn():
    """Build the serving inputs."""
    inputs = {}
    for feat in feature_columns:
        print(feat)
        inputs[feat.name] = tf.compat.v1.placeholder(shape=[None], dtype=feat.dtype)

    return tf.estimator.export.ServingInputReceiver(inputs, inputs)

path = est.export_saved_model('gs://tmp/btc',
                      json_serving_input_fn)

あるいは、不便ですが、以下の serving_input_fn を使って Proto Buffers に変換済みのバイトデータを入力するモデルを作ることもできます。


serving_input_fn = tf.estimator.export.build_parsing_serving_input_receiver_fn(
  tf.feature_column.make_parse_example_spec(feature_columns)
)

path = est.export_saved_model('gs://tmp/btc',
                      serving_input_fn)

BigQuery ML にする

CREATE OR REPLACE MODEL
  example_dataset.imported_tf_model OPTIONS (MODEL_TYPE='TENSORFLOW',
    MODEL_PATH='gs://tmp/btc/1581656284/*')

使ってみる

SELECT
  *
FROM
  ML.PREDICT(MODEL bigqueryml.gbdt,
    (
    SELECT
      *
    FROM
      UNNEST([STRUCT(10 AS age,
          0 AS fare), (80,
          30)])))

実行結果

行	all_class_ids	all_classes	class_ids	classes	logistic	logits	probabilities	age	fare
1	0	0	0	0	0.001847356558	-6.292150497	0.9981526732	10	0
	1	1					0.001847356558
2	0	0	1	1	0.9975745082	6.019282341	0.00242551649	80	30
	1	1					0.9975745082

いい感じに動いてそうです！

おわりに

BigQuery ML で使える TensorFlow の SavedModel を作って動作確認しました。
BigQuery ML でテンソルグラフ計算や、BigQuery ML では未リリースの BoostedTreesClassifier を実現できました。
BigQuery ML をうまく扱うと、データとモデルが近い位置におけるためぜひ活用していきたいですね。