如何从TFX BulkInferer获得数据帧或数据库写入

Question

我对TFX很陌生，但有一个显然可以工作的ML管道，它将通过BulkInferer使用。这似乎只产生Protobuf格式的输出，但由于我正在运行批量推理，我想将结果通过管道传输到数据库。(DB输出似乎应该是批量推理的默认输出，因为批量推理和DB访问都利用了并行化……但Protobuf是一种按记录序列化的格式。(

我想我可以使用Parquet Avro Protobuf之类的东西来进行转换(尽管它是在Java中，管道的其余部分是在Python中(，或者我可以自己写一些东西来逐个消费所有的Protobuf消息，将它们转换为JSON，将JSON反序列化为dict列表，并将dict加载到Pandas DataFrame中，或者将其存储为一堆键值对，我将其视为一次性数据库。。。但对于一个非常常见的用例来说，这听起来像是涉及并行化和优化的大量工作和痛苦。顶级Protobuf消息定义是Tensorflow的PredictionLog。

这个必须是一个常见的用例，因为像这样的TensorFlowModelAnalytics函数消耗Pandas DataFrames。我宁愿能够直接写入DB(最好是Google BigQuery(或Parquet文件(因为Parquet/Spark似乎比Pandas更好地并行化(，而且，这些似乎应该是常见的用例，但我还没有找到任何例子。也许我用错了搜索词？

我还查看了PredictExtractor；提取预测"；听起来很接近我想要的。。。但官方文档似乎对如何使用该类保持沉默。我认为TFTtransformOutput听起来像是一个很有前途的动词，但实际上它是一个名词。

我显然错过了一些基本的东西。没有人想将BulkInferer结果存储在数据库中是有原因的吗？是否有配置选项允许我将结果写入数据库？也许我想在TFX管道中添加一个ParquetIO或BigQueryIO实例？(TFX文档说它使用Beam"under the hood"，但这并没有说明我应该如何将它们一起使用。(但这些文档中的语法看起来与我的TFX代码有很大不同，我不确定它们是否兼容？

帮助？

Answer 1

(从相关问题复制以提高可见性(

经过一番挖掘，这里有一种替代方法，它假设事先不知道feature_spec。执行以下操作：

• 通过向组件构造添加output_example_spec，将BulkInferrer设置为写入output_examples而不是inference_result
• 在BulkInferrer之后的主管道中添加StatisticsGen和SchemaGen组件，以生成上述output_examples的模式
• 使用SchemaGen和BulkInferrer中的工件读取TFRecords并执行任何必要的操作

bulk_inferrer = BulkInferrer(
....
output_example_spec=bulk_inferrer_pb2.OutputExampleSpec(
output_columns_spec=[bulk_inferrer_pb2.OutputColumnsSpec(
predict_output=bulk_inferrer_pb2.PredictOutput(
output_columns=[bulk_inferrer_pb2.PredictOutputCol(
output_key='original_label_name',
output_column='output_label_column_name', )]))]
))
statistics = StatisticsGen(
examples=bulk_inferrer.outputs.output_examples
)
schema = SchemaGen(
statistics=statistics.outputs.output,
)

之后，可以执行以下操作：

import tensorflow as tf
from tfx.utils import io_utils
from tensorflow_transform.tf_metadata import schema_utils
# read schema from SchemaGen
schema_path = '/path/to/schemagen/schema.pbtxt'
schema_proto = io_utils.SchemaReader().read(schema_path)
spec = schema_utils.schema_as_feature_spec(schema_proto).feature_spec
# read inferred results
data_files = ['/path/to/bulkinferrer/output_examples/examples/examples-00000-of-00001.gz']
dataset = tf.data.TFRecordDataset(data_files, compression_type='GZIP')
# parse dataset with spec
def parse(raw_record):
return tf.io.parse_example(raw_record, spec)
dataset = dataset.map(parse)

在这一点上，该数据集与任何其他解析的数据集一样，因此编写CSV或BigQuery表或从中编写任何内容都是微不足道的。它确实帮助了我们在ZenML中的BatchInferencePipeline。

Answer 2

在这里回答我自己的问题，以记录我们所做的事情，尽管我认为@Hamza Tahir下面的回答客观上更好。这可能为需要更改开箱即用TFX组件的操作的其他情况提供了一个选项。不过它很有攻击性：

我们复制并编辑了文件tfx/components/bulk_deferrer/execute.py，替换了_run_model_inference()方法的内部管道中的转换：

| 'WritePredictionLogs' >> beam.io.WriteToTFRecord(
os.path.join(inference_result.uri, _PREDICTION_LOGS_FILE_NAME),
file_name_suffix='.gz',
coder=beam.coders.ProtoCoder(prediction_log_pb2.PredictionLog)))

这个：

| 'WritePredictionLogsBigquery' >> beam.io.WriteToBigQuery(
'our_project:namespace.TableName',
schema='SCHEMA_AUTODETECT',
write_disposition=beam.io.BigQueryDisposition.WRITE_APPEND,
create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED,
custom_gcs_temp_location='gs://our-storage-bucket/tmp',
temp_file_format='NEWLINE_DELIMITED_JSON',
ignore_insert_ids=True,
)

(这之所以有效，是因为当您导入BulkInferer组件时，每个节点的工作都会分配给运行在工作节点上的这些执行器，TFX会将自己的库复制到这些节点上。不过，它不会从用户空间库中复制所有，这就是为什么我们不能只对BulkInfrer进行子类化并导入自定义版本的原因。(

我们必须确保'our_project:namespace.TableName'的表具有与模型输出兼容的模式，但不必将该模式转换为JSON/AVRO。

理论上，我的团队希望使用围绕这一点构建的TFX来发出pull请求，但目前我们正在对几个关键参数进行硬编码，没有时间将其转化为真正的公共/生产状态。

Answer 3

我参加聚会有点晚了，但这是我用于此任务的一些代码：

import tensorflow as tf
from tensorflow_serving.apis import prediction_log_pb2
import pandas as pd

def parse_prediction_logs(inference_filenames: List[Text]): -> pd.DataFrame
"""
Args:
inference files:  tf.io.gfile.glob(Inferrer artifact uri)
Returns:
a dataframe of userids, predictions, and features
"""
def parse_log(pbuf):
# parse the protobuf
message = prediction_log_pb2.PredictionLog()
message.ParseFromString(pbuf)
# my model produces scores and classes and I extract the topK classes
predictions = [x.decode() for x in (message
.predict_log
.response
.outputs['output_2']
.string_val
)[:10]]
# here I parse the input tf.train.Example proto
inputs = tf.train.Example()
inputs.ParseFromString(message
.predict_log
.request
.inputs['input_1'].string_val[0]
)
# you can pull out individual features like this         
uid = inputs.features.feature["userId"].bytes_list.value[0].decode()
feature1 = [
x.decode() for x in inputs.features.feature["feature1"].bytes_list.value
]
feature2 = [
x.decode() for x in inputs.features.feature["feature2"].bytes_list.value
]
return (uid, predictions, feature1, feature2)
return pd.DataFrame(
[parse_log(x) for x in
tf.data.TFRecordDataset(inference_filenames, compression_type="GZIP").as_numpy_iterator()
], columns = ["userId", "predictions", "feature1", "feature2"]
)