TFXを使った機械学習パイプラインの構築（実装編その2）

システム開発

技術解説

連載 [第5回] :

徹底解説！Kubernetes上で動く機械学習プラットフォームKubeflow

2022年1月28日(金)

張替清音（はりがえきよなり）

連載の5回目となる今回は、前回に引き続きTFXを使ってKubeflow上で動かす機械学習パイプラインを構築していきます。

チューニング

ここまで、本連載の第3回で解説したモデルをベースにモデル学習のステップまで実装しました。しかし、性能の高いモデルを作成する上ではハイパーパラメータのチューニングが重要になります。ハイパーパラメータはモデルの種類により異なりますが、ディープラーニングモデルの場合は、主にオプティマイザーとネットワークアーキテクチャを制御するタイプのパラメータとなります。これらのパラメータを調整するプロセスは各パラメータの取り得る値の組み合わせを試行する必要があり、手動では骨が折れる作業です。TFXコンポーネントのTunerを使用しパイプラインに組み込むことで、このプロセスを自動化することができます。また、Tunerによって得られた結果を可視化し、最良のハイパーパラメータを特定するためのプロセスを支援するツールも提供されています。

Tunerとは

Tunerはハイパーパラメータチューニングを行うTFXコンポーネントです。執筆時点（2022年1月）ではKeras Tunerをライブラリとして利用することができます。※CloudTunerというGoogle Cloud Platform上で利用できるライブラリもありますが、こちらもバックエンドにKerasTunerを利用しています。TunerではTrainerと同様に主に次の入力を必要とします。

前回「データ変換」のステップで実装したデータ変換処理のグラフ
前回「データ検証」のステップで生成したデータスキーマ
前回「データ分割」のステップで説明した学習用データと評価用データ
学習パラメータ（学習ステップ、エポック数など）
ハイパーパラメータチューニングのプロセスを定義するtuner_fn関数を含むモジュールファイル

Tunerコンポーネントは与えられたデータおよびモデルを使用して、ハイパーパラメータを調整し最良の結果を出力します。

それでは、Tunerコンポーネントのエントリポイントである「tuner_fn関数」を実装します。なお、TunerもTrainerと同様にノートブック上で実行するとともにモジュールファイルとして書き出しを行います。

`1`	`tuner_module_file_name = 'adult_income_tuner_module.py'`

`2`	`tuner_module_file_path = os.path.join(os.getcwd(), tuner_module_file_name)`

次に「tuner_fn関数」を実装します。この関数では以下の処理を実装します。

学習用と評価用のデータの読み込み
利用するTunerの定義
チューニングの過程を記録するログの設定
結果の返却

クリックして表示

01%%writefile {tuner_module_file_path}
02 
03import kerastuner
04import tensorflow as tf
05import tensorflow_transform as tft
06 
07from tensorflow.keras import Model
08from tensorflow.keras.models import Sequential
09from tensorflow.keras.layers import Dense, Input, concatenate
10 
11from tfx.components.trainer.fn_args_utils import FnArgs
12from tfx.components.tuner.component import TunerFnResult
13 
14import adult_income_trainer_module_without_tuning as tm
15 
16def create_model(hparams: kerastuner.HyperParameters):
17    model = Sequential()
18 
19    inputs = []
20    for key in tm.NUMERIC_FEATURE_KEYS:
21        inputs.append(Input(shape=(1), name=tm.transformed_name(key)))
22    for key, dim in tm.ONE_HOT_FEATURES.items():
23        for i in range(0, dim):
24            inputs.append(Input(shape=(1), name=tm.transformed_name(key + '_' + str(i))))
25 
26    outputs = concatenate(inputs)
27    units = hparams.get('units')
28    outputs = Dense(units, activation='relu')(outputs)
29    outputs = Dense(units, activation='relu')(outputs)
30    outputs = Dense(1, activation='sigmoid')(outputs)
31 
32    model = Model(inputs=inputs, outputs=outputs)
33    model.compile(optimizer=hparams.get('optimizer'), loss='binary_crossentropy', metrics=['accuracy'])
34    model.summary()
35    return model
36 
37def tuner_fn(fn_args: FnArgs) -> TunerFnResult:
38    tuner = kerastuner.RandomSearch(
39        create_model,
40        max_trials = 10,
41        hyperparameters = get_hyperparameters(),
42        allow_new_entries = False,
43        objective = kerastuner.Objective('val_accuracy', 'max'),
44        directory = fn_args.working_dir,
45        project_name = 'census-income-tuner'
46    )
47 
48    tf_transform_graph = tft.TFTransformOutput(fn_args.transform_graph_path)
49    train_dataset = tm.input_fn(fn_args.train_files, fn_args.data_accessor,
50                                tf_transform_graph, 40)
51    eval_dataset = tm.input_fn(fn_args.eval_files, fn_args.data_accessor,
52                               tf_transform_graph, 40)
53 
54    tb_callback = tf.keras.callbacks.TensorBoard(
55        log_dir=fn_args.custom_config["tuner_log_dir"],
56        update_freq='batch')
57 
58    return TunerFnResult(
59        tuner = tuner,
60        fit_kwargs = {
61            'x': train_dataset,
62            'validation_data': eval_dataset,
63            'steps_per_epoch': fn_args.train_steps,
64            'validation_steps': fn_args.eval_steps,
65            'callbacks': [tb_callback]
66      })

Tunerは以下の3種類の探索アルゴリズムを選択可能です。上記の例では、RandomSearchを選択しています。

RandomSearch: 候補の値をランダムに組み合わせて探索します
BayesianOptimization: ベイズ最適化アルゴリズムを使い探索します
Hyperband: Hyperbandアルゴリズムを使い探索します

なお、目的とする評価指標は「val_accuracy」を指定しています。つまり、「評価用データに対する正解率」を最大化するハイパーパラメータを探索します。

次に「tuner_fn関数」に記述したヘルパー関数の「get_hyperparameters関数」を定義します。「get_hyperparameters関数」は探索するハイパーパラメータの値の種類、範囲を設定します。Keras TunerのHyperParametersクラスでは主に次のような設定ができます。

HyperParameters.Boolean: True or Falseを設定します
HyperParameters.Choice: 値の種類を設定します
HyperParameters.Int: 整数型の値の範囲を設定します。Stepを指定することで均等分割できます
HyperParameters.Float: 浮動小数点型の値の範囲を設定します。Stepを指定することで均等分割できます

ここでは、オプティマイザーとして「rmsprop」または「adam」、中間層のユニット数を100～250の範囲で50ずつ可変させながら探索します。

1%%writefile -a {tuner_module_file_path}
2 
3def get_hyperparameters() -> kerastuner.HyperParameters:
4    hp = kerastuner.HyperParameters()
5    hp.Choice('optimizer', ['rmsprop', 'adam'], default='adam')
6    hp.Int('units', min_value=100, max_value=250, step=50)
7    return hp

ハイパーパラメータチューニングを実行します。

01import tempfile
02 
03tuner_tensorboard_log_dir = tempfile.mkdtemp()
04param_tuner_custom_config_args = {"tuner_log_dir": tuner_tensorboard_log_dir}
05 
06param_tuner_train_args = {"num_steps": 100}
07param_tuner_eval_args = {"num_steps": 50}
08 
09tuner = tfx.components.Tuner(
10    examples = transform.outputs['transformed_examples'],
11    transform_graph = transform.outputs['transform_graph'],
12    schema = import_schema.outputs['result'],
13    module_file = tuner_module_file_path,
14    train_args = param_tuner_train_args,
15    eval_args = param_tuner_eval_args,
16    custom_config = param_tuner_custom_config_args
17)
18context.run(tuner)

図2-2：Tunerの実行結果

8回の試行を行い、最良のパラメータは次の通りになりました。

オプティマイザー： adam
中間層のユニット数： 150

次にTunerの結果を可視化して、詳細に解析してみます。可視化することにより、ハイパーパラメータが結果に与える影響を直感的に把握できるようになります。可視化にはTensorBoardというTensorFlowエコシステムの可視化ツールを利用します。

注釈

TensorBoardを利用するには、クライアント端末上で次のコマンドの実行および、設定が必要です。

クライアント端末で次のコマンドを実行
$ kubectl port-forward --address 0.0.0.0 example-notebook-0 6016:6016 -n anonymous
SSH クライアントのポートフォワーディングの設定に次を追加
ローカルのポート: 6016、リモート側ホスト：localhost、リモート側ポート: 6016

先ほど実施したTunerのログを読み込ませ、TensorBoardを起動します。

`1`	`%load_ext tensorboard`

`2`	`%tensorboard --logdir {tuner_tensorboard_log_dir} --host 0.0.0.0 --port 6016`

起動すると次のような画面がノートブックのセル内に表示されます。この画面では、各試行における損失曲線とメトリック曲線が表示されています。左ペインの操作パネルで可視化設定を切り替えることができます。

図2-3：TensorBoard-SCALARS画面

また、「GRAPHS」タブを選択すると、モデルグラフが表示されます。

図2-4：TensorBoard-GRAPHS画面

モデルの構造の概念グラフを可視化してくれるため、意図した設計と一致しているかどうかを確認できます。左ペインの「Run」の設定で各試行におけるモデルグラフを切り替えることができます。

ここまでは、TensorBoardの一般的な機能となります。これら機能に加えてハイパーパラメータチューニング時は、その試行を可視化する「HParams」が追加されます。画面上部の「HPARAMS」タブを選択します。

この画面では次のような項目を確認できます。

・TABLE VIEW

試行ごとに設定したハイパーパラメータ値と評価指標を確認できます。

図2-5：TABLE VIEW画面

左ペインのパネルで特定のハイパーパラメータのフィルタを指定できます。例えば、オプティマイザーが「rmsprop」、ユニット数が150以上の試行のみを表示できます。

図2-6：TABLE VIEWフィルタ結果画面

・PARALLEL COORDINATES VIEW

平行座標ビューでは、色付きの各線は試行を示しています。軸はハイパーパラメータと評価指標です。

図2-7：PARALLEL COORDINATES VIEW画面

TABLE VIEWと同様に左ペインのパネルでフィルタが可能です。指標値を見やすくするために、左ペインの「Metrics」の表示項目を「validation.epoch_accuracy」のみ表示するよう選択します。例えば、今回の結果を考察してみると、オプティマイザーが「adam」かつユニット数が150の試行が最も指標の値が高いことがわかります。また、オプティマイザーが「adam」の時の試行が指標の結果が良いことも視覚的にわかります。

図2-8：PARALLEL COORDINATES VIEWフィルタ結果画面

・SCATTER PLOT MATRIX VIEW

散布図マトリックスビューでは、各ドットは試行を表します。プロットは、軸としてさまざまなハイパーパラメータとメトリクスを使用した平面での試行の投影となっています。ハイパーパラメータごとに試行結果の分布を把握できます。

図2-9：SCATTER PLOT MATRIX VIEW画面

これら可視化した結果を確認して、ここではTunerコンポーネントを用いた試行で最良となった、オプティマイザー「adam」、中間層のユニット数「150」というハイパーパラメータを選定することにします。

・ハイパーパラメータを固定化して再度Trainerを実行

先ほど実施したTunerの結果で得られた最良のハイパーパラメータをモデル学習時に読み込めるようにImportTunerコンポーネントを使い、Trainerの入力として設定します。まず、固定化するパラメータ（Tunerの結果）をMinIOにアップロードします。

`1`	`bucket.upload_file(os.path.join(tuner.outputs['best_hyperparameters'].get()[0].uri, 'best_hyperparameters.txt'),`

`2`	`'tuner/best_hyperparameters.txt')`

ImportTunerコンポーネントを定義し、実行します。

1import_tuner = tfx.dsl.Importer(
2    source_uri = '%s/tuner/' % bucket_name_s3_prefix,
3    artifact_type = tfx.types.standard_artifacts.HyperParameters).with_id('tuner_importer')
4context.run(import_tuner)

図2-10：Import Tuner実行結果

TrainerコンポーネントでImportTunerを読み込むようにTrainerを再定義します。まずモジュールファイルを定義します。

`1`	`trainer_module_file_name = 'adult_income_trainer_module.py'`

`2`	`trainer_module_file_path = os.path.join(os.getcwd(), trainer_module_file_name)`

次に、チューニングして決定した最良のハイパーパラメータを受け取ってトレーニングをするために、Trainerコンポーネントで使用するモジュールファイルを更新します。

クリックして表示

001%%writefile {trainer_module_file_path}
002 
003import absl
004import os
005import sys
006 
007from typing import List, Text
008 
009import kerastuner
010import tensorflow as tf
011import tensorflow_transform as tft
012 
013from tensorflow.keras import Model
014from tensorflow.keras.models import Sequential
015from tensorflow.keras.layers import Dense, Input, concatenate
016 
017from tfx.components.trainer.fn_args_utils import DataAccessor
018from tfx.components.trainer.executor import TrainerFnArgs
019from tfx_bsl.tfxio import dataset_options
020 
021NUMERIC_FEATURE_KEYS = [
022    'age',
023    'education-num',
024    'capital-gain',
025    'capital-loss',
026    'hours-per-week',
027]
028ONE_HOT_FEATURES = {'workclass': 8,
029                    'education': 16,
030                    'marital-status': 7,
031                    'occupation': 14,
032                    'relationship': 6,
033                    'gender': 2
034                   }
035LABEL_KEY = 'income'
036 
037def transformed_name(key):
038    return key + '_xf'
039 
040def input_fn(file_pattern: List[Text],
041             data_accessor: DataAccessor,
042             tf_transform_output: tft.TFTransformOutput,
043             batch_size: int = 200) -> tf.data.Dataset:
044    return data_accessor.tf_dataset_factory(
045        file_pattern,
046        dataset_options.TensorFlowDatasetOptions(
047            batch_size=batch_size, label_key=transformed_name(LABEL_KEY)),
048        tf_transform_output.transformed_metadata.schema)
049 
050def create_model(hparams: kerastuner.HyperParameters):
051    model = Sequential()
052 
053    inputs = []
054    for key in NUMERIC_FEATURE_KEYS:
055        inputs.append(Input(shape=(1), name=transformed_name(key)))
056    for key, dim in ONE_HOT_FEATURES.items():
057        for i in range(0, dim):
058            inputs.append(Input(shape=(1), name=transformed_name(key + '_' + str(i))))
059 
060    outputs = concatenate(inputs)
061    units = hparams.get('units')
062    outputs = Dense(units, activation='relu')(outputs)
063    outputs = Dense(units, activation='relu')(outputs)
064    outputs = Dense(1, activation='sigmoid')(outputs)
065 
066    model = Model(inputs=inputs, outputs=outputs)
067    model.compile(optimizer=hparams.get('optimizer'), loss='binary_crossentropy', metrics=['accuracy'])
068    model.summary()
069    return model
070 
071def get_tf_examples_serving_signature(model, tf_transform_output):
072    model.tft_layer_inference = tf_transform_output.transform_features_layer()
073 
074    @tf.function(input_signature=[
075        tf.TensorSpec(shape=[None], dtype=tf.string, name='examples')
076    ])
077    def serve_tf_examples_fn(serialized_tf_example):
078        raw_feature_spec = tf_transform_output.raw_feature_spec()
079        raw_feature_spec.pop(LABEL_KEY)
080        raw_features = tf.io.parse_example(serialized_tf_example, raw_feature_spec)
081        transformed_features = model.tft_layer_inference(raw_features)
082 
083        outputs = model(transformed_features)
084        return {'outputs': outputs}
085 
086    return serve_tf_examples_fn
087 
088def get_transform_features_signature(model, tf_transform_output):
089    model.tft_layer_eval = tf_transform_output.transform_features_layer()
090 
091    @tf.function(input_signature=[
092        tf.TensorSpec(shape=[None], dtype=tf.string, name='examples')
093    ])
094    def transform_features_fn(serialized_tf_example):
095        raw_feature_spec = tf_transform_output.raw_feature_spec()
096        raw_features = tf.io.parse_example(serialized_tf_example, raw_feature_spec)
097        transformed_features = model.tft_layer_eval(raw_features)
098        return transformed_features
099 
100    return transform_features_fn
101 
102def run_fn(fn_args: TrainerFnArgs):
103    tf_transform_output = tft.TFTransformOutput(fn_args.transform_output)
104    train_dataset = input_fn(fn_args.train_files, fn_args.data_accessor,
105                                tf_transform_output, 40)
106    eval_dataset = input_fn(fn_args.eval_files, fn_args.data_accessor,
107                               tf_transform_output, 40)
108 
109    hparams = kerastuner.HyperParameters.from_config(fn_args.hyperparameters)
110    absl.logging.info('HyperParameters for training: %s' % hparams.get_config())
111 
112    tensorboard_callback = tf.keras.callbacks.TensorBoard(
113        log_dir = fn_args.model_run_dir, update_freq='batch'
114    )
115 
116    model = create_model(hparams)
117    model.fit(
118        train_dataset,
119        epochs = fn_args.custom_config["epoch"],
120        steps_per_epoch = fn_args.train_steps,
121        validation_data = eval_dataset,
122        validation_steps = fn_args.eval_steps,
123        callbacks = [tensorboard_callback]
124    )
125 
126    signatures = {
127        'serving_default':
128            get_tf_examples_serving_signature(model, tf_transform_output),
129        'transform_features':
130            get_transform_features_signature(model, tf_transform_output),
131    }
132    model.save(fn_args.serving_model_dir, save_format='tf', signatures=signatures)

モジュールファイルをMinIOにアップロードします。

`1`	`bucket.upload_file(trainer_module_file_path, 'module/%s' % trainer_module_file_name)`

Trainerの入力としてhyperparametersを追加し、先ほど作成したモジュールファイルを使用してモデル学習を再実行します。

01param_trainer_module_file = '%s/module/%s' % (bucket_name_s3_prefix, trainer_module_file_name)
02 
03trainer = tfx.components.Trainer(
04    module_file = param_trainer_module_file,
05    custom_executor_spec = executor_spec.ExecutorClassSpec(GenericExecutor),
06    examples = transform.outputs['transformed_examples'],
07    transform_graph = transform.outputs['transform_graph'],
08    schema = import_schema.outputs['result'],
09    hyperparameters = import_tuner.outputs['result'],
10    train_args = param_trainer_train_args,
11    eval_args = param_trainer_eval_args,
12    custom_config = param_trainer_custom_config_args)
13context.run(trainer)