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

transform_module_file_name = 'adult_income_transorm_module.py'
transform_module_file_path = os.path.join(os.getcwd(), transform_module_file_name)

%%writefile {transform_module_file_path}
import os
import sys

import tensorflow as tf
import tensorflow_transform as tft

# 数値変数の項目名リスト
NUMERIC_FEATURE_KEYS = [
    'age',
    'education-num',
    'capital-gain',
    'capital-loss',
    'hours-per-week',
]
# カテゴリー変数の項目名とその次元数の辞書
ONE_HOT_FEATURES = {'workclass': 8,
                    'education': 16,
                    'marital-status': 7,
                    'occupation': 14,
                    'relationship': 6,
                    'gender': 2
                   }
# 予測対象の項目名
LABEL_KEY = 'income'

# 変換後の項目名の変換用関数
def transformed_name(key):
    return key + '_xf'

%%writefile -a {transform_module_file_path}

def preprocessing_fn(inputs):
    #文字列型の値を数値型に変換する関数
    def __convert_string2int_value(key, original_values, converted_values):
        initializer = tf.lookup.KeyValueTensorInitializer(
            keys=original_values,
            values=tf.cast(converted_values, tf.int64),
            key_dtype=tf.string,
            value_dtype=tf.int64)
        table = tf.lookup.StaticHashTable(initializer, default_value=-1)
        return table.lookup(inputs[key])

    outputs = {}
    #欠損値を最頻値に変換
    inputs['workclass'] = tf.where(tf.equal(inputs['workclass'], '?'),
                                            'Private', inputs['workclass'])
    inputs['occupation'] = tf.where(tf.equal(inputs['occupation'], '?'),
                                             'Prof-specialty', inputs['occupation'])

    # TFTの関数を使って数値型の項目を正規化 (平均値0標準偏差1に変換)
    for key in NUMERIC_FEATURE_KEYS:
        outputs[transformed_name(key)] = tft.scale_to_z_score(inputs[key])

    # TFT の関数を使ってダミー変数化(One-Hot エンコーディング)
    for key in ONE_HOT_FEATURES.keys():
        dim = ONE_HOT_FEATURES[key]
        index = tft.compute_and_apply_vocabulary(
            tf.strings.strip(inputs[key]),
            num_oov_buckets=0,
            vocab_filename=key)
        one_hot_tensor = tf.one_hot(index, dim)
        one_hot_tensor = tf.reshape(one_hot_tensor, [-1, dim])
        for i in range(0, dim):
            one_hot_name = transformed_name(key + '_' + str(i))
            outputs[one_hot_name] = one_hot_tensor[:,i]
    # Labelを0と1に変換
    table_keys = ['<=50K', '>50K']
    outputs[transformed_name(LABEL_KEY)] = __convert_string2int_value(
             LABEL_KEY,  table_keys, [0, 1])

    return outputs

bucket.upload_file(transform_module_file_path, 'module/%s' % transform_module_file_name)

param_transform_module_file = '%s/module/%s' % (bucket_name_s3_prefix, transform_module_file_name)

transform = tfx.components.Transform(
    examples = example_gen.outputs['examples'],
    schema = import_schema.outputs['result'],
    module_file = param_transform_module_file)
context.run(transform)

import_anomalies_schema = tfx.dsl.Importer(
    source_uri = update_schema_output_dir,
    artifact_type = tfx.types.standard_artifacts.Schema).with_id('schema_importer')
context.run(import_anomalies_schema)

anomalies_example_validator = tfx.components.ExampleValidator(
    statistics = statistics_gen.outputs['statistics'],
    schema = import_anomalies_schema.outputs['result'])
context.run(anomalies_example_validator)

context.show(anomalies_example_validator.outputs['anomalies'])

図2-17：ExampleValidatorで異常値検出した実行結果

ageに有効範囲外の数値が含まれているのと、workclassに取り得る値ではない「?」という値が含まれているため、異常として検出されました。この例は意図的に実際のデータセットの構造ではないスキーマを定義し、異常を検出させましたが、新しく取り込むデータセットの項目数や取り得る値の種類、値の範囲が変わった場合などに、この例と同様に異常検出することができます。

今回、パイプラインで利用する「正しいスキーマ」は、前述した「データ分析」の「スキーマ生成」でアップロードしたスキーマとなるため、それを読み込むExampleValidatorを定義して再度データ検証を実施します。

example_validator = tfx.components.ExampleValidator(
    statistics = statistics_gen.outputs['statistics'],
    schema = import_schema.outputs['result'])
context.run(example_validator)

図2-18：ExampleValidatorの実行結果

データセットの検証結果を表示します。

context.show(example_validator.outputs['anomalies'])

図2-19：ExampleValidatorで異常値が存在しなかった場合の実行結果

異常は検出されませんでした。

• データ変換済みのデータに対するデータ検証

Transformコンポーネントには標準でデータ検証のコンポーネントが組み込まれているため、それを利用する方法について解説します。先ほど実行したTransformの結果からデータ変換後のスキーマを参照します。

context.show(transform.outputs['post_transform_schema'])

図2-20：Transform実行後のスキーマ表示の一部

このスキーマで問題ないと仮定して、そのスキーマをパイプラインで利用するために、MinIOにアップロードします。

bucket.upload_file(os.path.join(transform.outputs['post_transform_schema'].get()[0].uri, 'schema.pbtxt'), 'post_schema/schema.pbtxt')

次にTransformコンポーネントで利用するスキーマを設定します。stats_options_updater_fn関数を追加することで、既存のスキーマをTransformコンポーネントに読み込ませることができます。

%%writefile -a {transform_module_file_path}

def stats_options_updater_fn(stats_type, stats_options):
    import tensorflow_data_validation as tfdv
    from tfx.components.transform import stats_options_util

    if stats_type == stats_options_util.StatsType.POST_TRANSFORM:
        load_post_schema = tfdv.load_schema_text('s3://census-income/post_schema/schema.pbtxt')
        stats_options.schema = load_post_schema
    return stats_options

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

bucket.upload_file(transform_module_file_path, 'module/%s' % transform_module_file_name)

これで、入力データセットに対するデータ検証と変換済みデータセットに対するデータ検証をパイプラインに組み込むことができました。

データ分割（分割済みデータの取り出し）

ここまで入力データの取り込みからデータを変換し検証するところまでパイプラインの実装を進めてきました。本稿の最後に機械学習モデルの学習時に利用する「学習用データ」と「評価用データ」を取り出し、その内容をノートブック上で確認する手順を解説します。※TFXパイプラインでは、データセットの分割をExampleGenで行います。分割に関する設定は前述した「入力データの取り込み」の「データセットの分割設定」を参照してください。

• 学習用データと評価用データの内容を確認

学習用データは「Split-train」を、評価用は「Split-eval」をそれぞれキーとして指定することで、Transformコンポーネントの出力結果から変換済みのデータファイルのパスを取得できます。次の例は学習用のデータを読み込み、表示する手順です。

transformed_train_uri = os.path.join(transform.outputs['transformed_examples'].get()[0].uri, 'Split-train')
tfrecord_filenames = [os.path.join(transformed_train_uri, name)
                      for name in os.listdir(transformed_train_uri)]

dataset = tf.data.TFRecordDataset(tfrecord_filenames, compression_type="GZIP")
#上から3行を表示
for tfrecord in dataset.take(3):
    serialized_example = tfrecord.numpy()
    example = tf.train.Example()
    example.ParseFromString(serialized_example)
    print(example)

図2-21：Transform実行後の「Split-train」のデータ表示の一部

また、Transformコンポーネントの出力結果には、統計情報も含まれています。これを取得してTFDVで可視化することも可能です。

post_transform_path = os.path.join(transform.outputs['post_transform_stats'].get()[0].uri, 'FeatureStats.pb')
post_stats = tfdv.load_stats_binary(post_transform_path)
tfdv.visualize_statistics(post_stats)

図2-22：TFDVによる学習用データの可視化

ここまで、機械学習パイプラインのうち「データ準備」を構成する各ステップについて、TFXコンポーネントを使って実装しました。

おわりに

今回はTFXの概要と機械学習パイプラインのうち「データ準備」のステップについて、TFXコンポーネントの解説と実装を行いました。次回は、後半部分となる「モデル作成」について、引き続き解説と実装を行います。