BigQuery ML で使える TensorFlow モデルを作る

# はじめに
BigQuery ML は [インポートした TensorFlow モデルでの予測](https://cloud.google.com/bigquery-ml/docs/making-predictions-with-imported-tensorflow-models) ができます。
BigQuery ML を使って TensorFlow モデルを管理すると、データの転送を省略したり、モデルを BigQuery と Cloud Storage に一任できます。
この記事では、BigQuery ML にインポート可能な TensorFlow モデルを作っていきます。

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


# モデルの作り方
[TensorFlow モデルをインポートする CREATE MODEL ステートメント](https://cloud.google.com/bigquery-ml/docs/reference/standard-sql/bigqueryml-syntax-create-tensorflow) にあるように、BigQuery ML で使える TensorFlow モデルは SavedModel として保存されている必要があります。

# SavedModel を作る

## シンプルな SavedModel を作る
[tf.saved_model.save](https://www.tensorflow.org/api_docs/python/tf/saved_model/save) の例から始めましょう。
例では、tf.TensorSpec は `shape=None` と定義されていますが、BigQuery ML から使う場合は必須のようですので、
`shape=1` とします。関数を変える場合には、適切に tf.TensorSpec を書き換えてください。

```python
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 モデルのインポート](https://cloud.google.com/bigquery-ml/docs/making-predictions-with-imported-tensorflow-models#importing_models) を参考にモデルをインポートします。
Cloud Storage にある SavedModel を参照できるので、予め転送しておきましょう。
クエリ 1 つでインポートできるのでとてもお手軽です。

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

### インポートしたモデルを使う
[インポートした TensorFlow モデルでの予測](https://cloud.google.com/bigquery-ml/docs/making-predictions-with-imported-tensorflow-models#making_predictions_with_imported_models) を参考にモデルで予測します。

```sql
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](https://www.tensorflow.org/guide/autograph) によって、Python コードが TensorFlow グラフに変換可能です。JavaScript UDF の代わりに使えるかもしれません。

## サポートされている型
[サポートされている入力](https://cloud.google.com/bigquery-ml/docs/reference/standard-sql/bigqueryml-syntax-create-tensorflow#inputs) にありますが、再掲します。

|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 のエクスポート](https://cloud.google.com/ml-engine/docs/tensorflow/exporting-for-prediction) を参考に、機械学習モデルの SavedModel を作っていきます。

ここでは、BigQuery ML のリリースが待たれる Boosted Trees として、[Boosted trees using Estimators](https://www.tensorflow.org/tutorials/estimator/boosted_trees) を見ながら、tf.estimator.BoostedTreesClassifier を作ってみましょう。

### データのロード
```python
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')
```

### 入力値の作成
```python
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 の作成
```python
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 が使えるのでこのコードが動作します。

```python
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 氏より）。

```python

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 に変換済みのバイトデータを入力するモデルを作ることもできます。

```python

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 にする
```sql
CREATE OR REPLACE MODEL
  example_dataset.imported_tf_model OPTIONS (MODEL_TYPE='TENSORFLOW',
    MODEL_PATH='gs://tmp/btc/1581656284/*')
```

### 使ってみる
```sql
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 をうまく扱うと、データとモデルが近い位置におけるためぜひ活用していきたいですね。