5.6 TensorFlow Extended (TFX) 机器学习流水线 5.6 TensorFlow Extended (TFX) 机器学习流水线 TensorFlow Extended (TFX) 是一个用于部署生产级机器学习流水线的端到端平台。它提供了一套强大的工具和组件,用于数据验证、特征工程、模型训练、评估和部署,确保机器学习模型在生产环境中可靠、高效地运行。 5.6.1 TFX 的核心概念 流水线 (Pipeline): 定义了机器学习工作流的各个阶段,包括数据摄取、预处理、模型训练、评估和部署。
TensorFlow Extended (TFX) 是一个用于部署生产级机器学习流水线的端到端平台。它提供了一套强大的工具和组件,用于数据验证、特征工程、模型训练、评估和部署,确保机器学习模型在生产环境中可靠、高效地运行。
流水线 (Pipeline): 定义了机器学习工作流的各个阶段,包括数据摄取、预处理、模型训练、评估和部署。
组件 (Component): 流水线中的一个独立任务单元,例如 ExampleGen (数据摄取)、StatisticsGen (统计信息生成)、SchemaGen (模式推断)、Trainer (模型训练) 等。
工件 (Artifact): 组件的输入和输出,例如数据集、统计信息、模式、训练好的模型等。工件在流水线中流动,连接不同的组件。
元数据 (Metadata): 记录了流水线的执行历史、组件的输入输出、工件的版本信息等。TFX 使用元数据跟踪 lineage,方便调试和回溯。
编排器 (Orchestrator): 负责调度和执行流水线中的组件。TFX 支持多种编排器,例如 Kubeflow Pipelines、Apache Airflow 和 Beam。
一个典型的 TFX 流水线通常包含以下步骤:
数据摄取 (ExampleGen): 从各种数据源(例如 CSV 文件、数据库、云存储)读取数据,并将其转换为 TFX 使用的 tf.Example 格式。
统计信息生成 (StatisticsGen): 计算数据集的统计信息,例如均值、方差、最大值、最小值等。
模式推断 (SchemaGen): 根据统计信息推断数据集的模式 (Schema),包括每个特征的数据类型、取值范围、是否缺失等。
数据验证 (ExampleValidator): 根据模式验证数据集,检测异常值、缺失值、数据类型错误等。
特征工程 (Transform): 对数据进行预处理和特征工程,例如归一化、标准化、独热编码、特征交叉等。
模型训练 (Trainer): 使用预处理后的数据训练机器学习模型。
模型评估 (Evaluator): 评估训练好的模型在验证集上的性能指标,例如准确率、召回率、F1 值等。
模型验证 (InfraValidator): 在真实的部署环境中验证模型是否能够正常运行。
模型推送 (Pusher): 将验证通过的模型推送到模型服务器,例如 TensorFlow Serving。
以下是一个使用 TFX 构建简单流水线的示例,用于训练一个简单的线性回归模型。
1. 数据准备
首先,我们需要准备一些训练数据。这里我们使用一个简单的 CSV 文件,包含两个特征 feature_1 和 feature_2,以及一个标签 label。
feature_1,feature_2,label 1.0,2.0,3.0 2.0,3.0,5.0 3.0,4.0,7.0 4.0,5.0,9.0 5.0,6.0,11.0
2. 定义流水线
接下来,我们定义 TFX 流水线。首先,我们需要导入必要的库:
import os from tfx.components import CsvExampleGen from tfx.components import StatisticsGen from tfx.components import SchemaGen from tfx.components import ExampleValidator from tfx.components import Transform from tfx.components import Trainer from tfx.components import Evaluator from tfx.components import Pusher from tfx.orchestration import pipeline from tfx.orchestration.beam.beam_dag_runner import BeamDagRunner from tfx.proto import example_gen_pb2 from tfx.proto import trainer_pb2 import tensorflow as tf
然后,我们定义流水线的参数:
PIPELINE_NAME = 'my_pipeline' PIPELINE_ROOT = os.path.join('.', 'pipeline') DATA_ROOT = os.path.join('.', 'data') MODULE_FILE = os.path.join('.', 'model.py') SERVING_MODEL_DIR = os.path.join('.', 'serving_model') # 创建目录 os.makedirs(PIPELINE_ROOT, exist_ok=True) os.makedirs(DATA_ROOT, exist_ok=True) os.makedirs(SERVING_MODEL_DIR, exist_ok=True) # 将 CSV 数据写入文件 data_filepath = os.path.join(DATA_ROOT, 'data.csv') with open(data_filepath, 'w') as f: f.write("feature_1,feature_2,label\n") f.write("1.0,2.0,3.0\n") f.write("2.0,3.0,5.0\n") f.write("3.0,4.0,7.0\n") f.write("4.0,5.0,9.0\n") f.write("5.0,6.0,11.0\n")
接下来,我们定义各个组件:
# ExampleGen example_gen = CsvExampleGen(input_base=DATA_ROOT) # StatisticsGen statistics_gen = StatisticsGen(examples=example_gen.outputs['examples']) # SchemaGen schema_gen = SchemaGen( statistics=statistics_gen.outputs['statistics'], infer_feature_shape=False) # ExampleValidator example_validator = ExampleValidator( statistics=statistics_gen.outputs['statistics'], schema=schema_gen.outputs['schema']) # Transform transform = Transform( examples=example_gen.outputs['examples'], schema=schema_gen.outputs['schema'], module_file=MODULE_FILE) # Trainer trainer = Trainer( module_file=MODULE_FILE, examples=transform.outputs['transformed_examples'], transform_graph=transform.outputs['transform_graph'], schema=schema_gen.outputs['schema'], train_args=trainer_pb2.TrainArgs(num_steps=1000), eval_args=trainer_pb2.EvalArgs(num_steps=500)) # Evaluator evaluator = Evaluator( examples=example_gen.outputs['examples'], model=trainer.outputs['model']) # Pusher pusher = Pusher( model=trainer.outputs['model'], model_blessing=evaluator.outputs['blessing'], push_destination=pusher_pb2.PushDestination( filesystem=pusher_pb2.PushDestination.Filesystem( base_directory=SERVING_MODEL_DIR)))
最后,我们将所有组件组合成一个流水线:
components = [ example_gen, statistics_gen, schema_gen, example_validator, transform, trainer, evaluator, pusher, ] pipeline_def = pipeline.Pipeline( pipeline_name=PIPELINE_NAME, pipeline_root=PIPELINE_ROOT, components=components, enable_cache=True, metadata_connection_config=sqlite_metadata_connection_config( os.path.join(PIPELINE_ROOT, 'metadata.db')), )
3. 定义 model.py
model.py 文件包含了特征工程和模型定义的代码。以下是一个简单的示例:
import tensorflow as tf from tfx.components.trainer.fn_args_utils import FnArgs from tfx.components.transform.executor import Executor as TransformExecutor from tfx.orchestration.experimental.feature_utils import StandardFeature from tfx.proto import example_gen_pb2 from tfx.utils import schema_utils _FEATURE_KEYS = ['feature_1', 'feature_2'] _LABEL_KEY = 'label' def _input_fn(file_pattern: str, tf_transform_output: TransformExecutor.TransformOutput, batch_size: int = 32) -> tf.data.Dataset: """Creates a tf.data dataset for training or evaluation.""" transformed_feature_spec = ( tf_transform_output.transformed_feature_spec().copy()) dataset = tf.data.experimental.make_batched_features_dataset( file_pattern=file_pattern, batch_size=batch_size, features=transformed_feature_spec, reader=tf.data.TFRecordDataset, label_key=_LABEL_KEY) return dataset def _build_keras_model(tf_transform_output: TransformExecutor.TransformOutput) -> tf.keras.Model: """Creates a Keras model for regression.""" feature_spec = tf_transform_output.transformed_feature_spec() input_layers = { feature_name: tf.keras.layers.Input(shape=(1,), name=feature_name) for feature_name in feature_spec } concat_layer = tf.keras.layers.Concatenate()(list(input_layers.values())) dense_1 = tf.keras.layers.Dense(10, activation='relu')(concat_layer) output_layer = tf.keras.layers.Dense(1)(dense_1) model = tf.keras.Model(inputs=input_layers, outputs=output_layer) model.compile(optimizer='adam', loss='mse', metrics=['mae']) return model def run_fn(fn_args: FnArgs): """Defines the training loop.""" tf_transform_output = TransformExecutor.TransformOutput( fn_args.transform_output) train_dataset = _input_fn( fn_args.train_files[0], tf_transform_output, batch_size=fn_args.train_batch_size) eval_dataset = _input_fn( fn_args.eval_files[0], tf_transform_output, batch_size=fn_args.eval_batch_size) model = _build_keras_model(tf_transform_output) model.fit( train_dataset, steps_per_epoch=fn_args.train_steps_per_epoch, validation_data=eval_dataset, validation_steps=fn_args.eval_steps_per_epoch) model.save(fn_args.serving_model_dir, save_format='tf') def transform_fn(inputs): """tf.transform's callback function. Args: inputs: map from feature keys to `Tensor` of the same shape. Returns: A map from transformed feature names to the `Tensor`-s. """ outputs = {} for feature_name in _FEATURE_KEYS: outputs[feature_name] = tf.identity(inputs[feature_name]) # No transformation is applied in this example. return outputs
4. 运行流水线
最后,我们使用 BeamDagRunner 运行流水线:
from tfx.orchestration.sqlite.sqlite_metadata_db import SqliteMetadataSource from ml_metadata.proto import metadata_store_pb2 def sqlite_metadata_connection_config(sqlite_file: str) -> metadata_store_pb2.ConnectionConfig: """Create a connection config for an SQLite database.""" return metadata_store_pb2.ConnectionConfig( sqlite=metadata_store_pb2.SqliteMetadataSourceConfig( filename_uri=sqlite_file ), artifact_purge_mode=metadata_store_pb2.ConnectionConfig.ArtifactPurgeMode.DISABLED ) BeamDagRunner().run(pipeline_def)
自动化: TFX 自动化了机器学习流水线的各个阶段,减少了人工干预,提高了效率。
可靠性: TFX 提供了数据验证、模型评估和模型验证等机制,确保模型在生产环境中可靠运行。
可扩展性: TFX 可以轻松地扩展到大规模数据集和复杂的模型。
可追溯性: TFX 使用元数据跟踪 lineage,方便调试和回溯。
可重用性: TFX 组件可以被重用在不同的流水线中,提高了开发效率。
TensorFlow Extended (TFX) 是一个强大的平台,用于构建和部署生产级机器学习流水线。它提供了一套完整的工具和组件,可以自动化机器学习工作流的各个阶段,确保模型在生产环境中可靠、高效地运行。通过学习和使用 TFX,您可以更好地管理和维护您的机器学习模型,并从中获得更大的价值。