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

import os
import pandas as pd
import tensorflow as tf
import tensorflow_data_validation as tfdv
import tensorflow_transform as tft
import tensorflow_model_analysis as tfma

from tfx import v1 as tfx
from tfx.components.trainer.executor import GenericExecutor
from tfx.dsl.components.base import executor_spec
from tfx.orchestration.experimental.interactive.interactive_context import InteractiveContext

# TensorFlowの警告出力を抑止する
import logging
import warnings
logger = tf.get_logger()
logger.setLevel(logging.ERROR)
warnings.filterwarnings('ignore')

os.environ['UNSUPPORTED_DO_NOT_PACKAGE_USER_MODULES'] = 'True'

context = InteractiveContext()

header = [
    'age','workclass','id','education','education-num','marital-status',
    'occupation','relationship','race','gender','capital-gain','capital-loss',
    'hours-per-week','native-country','income'
]

df = pd.read_csv('adult.data', names=header, delimiter=', ', usecols=lambda x: x not in ['id', 'race', 'native-country'])

def _bytes_feature(value):
    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))

def _int64_feature(value):
    return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))

def create_tf_example(features, label):
    tf_example = tf.train.Example(features=tf.train.Features(feature={
        'age': _int64_feature(features[0]),
        'workclass': _bytes_feature(features[1].encode('utf-8')),
        'education': _bytes_feature(features[2].encode('utf-8')),
        'education-num': _int64_feature(features[3]),
        'marital-status': _bytes_feature(features[4].encode('utf-8')),
        'occupation': _bytes_feature(features[5].encode('utf-8')),
        'relationship': _bytes_feature(features[6].encode('utf-8')),
        'gender': _bytes_feature(features[7].encode('utf-8')),
        'capital-gain': _int64_feature(features[8]),
        'capital-loss': _int64_feature(features[9]),
        'hours-per-week': _int64_feature(features[10]),
        'income': _bytes_feature(label.encode('utf-8')),
    }))
    return tf_example

tfrecord_data_root = os.path.join(os.getcwd(), 'tfrecord_data', 'span-1', 'v1')
tfrecord_data_path = os.path.join(tfrecord_data_root, 'adult.span-1.tfrecords')
os.makedirs(tfrecord_data_root, exist_ok=True)

with tf.io.TFRecordWriter(tfrecord_data_path) as writer:
    for row in df.values:
        features, label = row[:-1], row[-1]
        example = create_tf_example(features, label)
        writer.write(example.SerializeToString())

import boto3
from botocore.client import Config

os.environ['AWS_ACCESS_KEY_ID'] = 'minio'
os.environ['AWS_SECRET_ACCESS_KEY'] = 'minio123'
os.environ["AWS_REGION"] = 'ap-northeast-1'
os.environ["S3_ENDPOINT"] = "http://minio-service.kubeflow.svc.cluster.local:9000"

s3_resource = boto3.resource('s3',
                endpoint_url = os.environ["S3_ENDPOINT"],
                aws_access_key_id = os.environ['AWS_ACCESS_KEY_ID'],
                aws_secret_access_key = os.environ['AWS_SECRET_ACCESS_KEY'],
                region_name = os.environ["AWS_REGION"],
                config = Config(signature_version='s3v4'))

bucket_name = 'census-income'
bucket_name_s3_prefix = 's3://%s' % bucket_name

try:
    bucket = s3_resource.Bucket(bucket_name)
    bucket.create()
    print("Create S3 bucket:", bucket_name)
except:
    print("Exsiting S3 bucket:", bucket_name)

bucket.upload_file(tfrecord_data_path, 'data/span-1/v1/adult.span-1.tfrecords')

beam_pipeline_args = ['--s3_access_key_id=%s' % os.environ['AWS_ACCESS_KEY_ID'],
                '--s3_secret_access_key=%s' % os.environ['AWS_SECRET_ACCESS_KEY'],
                '--s3_endpoint_url=%s' % os.environ["S3_ENDPOINT"],
                '--s3_region_name=%s' % os.environ["AWS_REGION"],
                '--s3_verify=0',
                '--s3_disable_ssl',
                '--direct_num_workers=1']

context.beam_pipeline_args=beam_pipeline_args

example_gen_input = {"splits": [{"name": "single_split",
                "pattern": "span-{SPAN}/v{VERSION}/*"}]
                    }
example_gen_output = {"split_config": {"splits": [{"name": "train", "hash_buckets": 2},
                {"name": "eval", "hash_buckets": 1}]
                                      }}
example_gen = tfx.components.ImportExampleGen(input_base = '%s/data/' % bucket_name_s3_prefix,
                input_config = example_gen_input,
                output_config = example_gen_output)

context.run(example_gen)

example_gen_input = {"splits": [{"name": "single_split",
                                 "pattern": "span-{SPAN}/v{VERSION}/*"}]
                    }

example_gen_output_three_division = {"split_config": {"splits": [{"name": "train", "hash_buckets": 6},
                                {"name": "eval", "hash_buckets": 2},
                                {"name": "test", "hash_buckets": 2}]}}
example_gen_three_division = tfx.components.ImportExampleGen(input_base = '%s/data/' % bucket_name_s3_prefix,
                                input_config = example_gen_input,
                                output_config = example_gen_output_three_division)
context.run(example_gen_three_division)