如何從 TFX BulkInferrer 獲取 dataframe 或數據庫寫入？

Question

我對 TFX 很陌生，但有一個顯然可以通過BulkInferrer使用的 ML 管道。 這似乎僅以 Protobuf 格式生成 output ，但由於我正在運行批量推理，我想將 pipe 的結果存儲到數據庫中。 （DB output 似乎應該是批量推理的默認值，因為批量推理和數據庫訪問都利用了並行化......但 Protobuf 是每記錄的序列化格式。）

我假設我可以使用Parquet-Avro-Protobuf之類的東西來進行轉換（盡管這是在 Java 和 Python 中管道的 rest 中），或者我可以自己編寫一些東西來一個接一個地使用所有 protobuf 消息， them into JSON, deserialize the JSON into a list of dicts, and load the dict into a Pandas DataFrame, or store it as a bunch of key-value pairs which I treat like a single-use DB... but that sounds like a對於一個非常常見的用例，涉及並行化和優化的大量工作和痛苦。 頂級 Protobuf 消息定義是 Tensorflow 的PredictionLog 。

這一定是一個常見的用例，因為像這樣的 TensorFlowModelAnalytics 函數會消耗 Pandas DataFrames。 我寧願能夠直接寫入數據庫（最好是 Google BigQuery）或 Parquet 文件（因為 Parquet / Spark 似乎比 Pandas 並行化更好），而且這些似乎應該是常見的用例，但我沒有找到任何例子。 也許我使用了錯誤的搜索詞？

我還查看了PredictExtractor ，因為“提取預測”聽起來很接近我想要的……但官方文檔似乎對 class 應該如何使用保持沉默。 我認為TFTransformOutput聽起來像是一個很有前途的動詞，但實際上它是一個名詞。

我顯然在這里遺漏了一些基本的東西。 沒有人願意將 BulkInferrer 結果存儲在數據庫中嗎？ 是否有允許我將結果寫入數據庫的配置選項？ 也許我想將ParquetIO或BigQueryIO實例添加到 TFX 管道？ （TFX 文檔說它在“幕后”使用 Beam，但這並沒有說明我應該如何一起使用它們。）但是這些文檔中的語法看起來與我的 TFX 代碼完全不同，我不確定它們是否'重新兼容？

幫助？

Answer 1

（從相關問題復制以提高知名度）

經過一番挖掘，這是一種替代方法，它假設事先不了解feature_spec 。 請執行下列操作：

• 通過將output_example_spec添加到組件構造中，將BulkInferrer設置為寫入output_examples而不是inference_result 。
• 在 BulkInferrer 之后的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 表或從那里的任何東西都是微不足道的。 它確實通過BatchInferencePipeline在ZenML中幫助了我們。

Answer 2

在這里回答我自己的問題以記錄我們所做的事情，盡管我認為下面@Hamza Tahir 的回答客觀上更好。 這可以為需要更改開箱即用 TFX 組件的操作的其他情況提供一個選項。 雖然它很hacky：

我們復制並編輯了文件 tfx/components/bulk_inferrer/executor.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,
       )

（這是有效的，因為當您導入 BulkInferrer 組件時，每個節點的工作被外包給在工作節點上運行的這些執行程序，並且 TFX 將其自己的庫復制到這些節點上。它不會從用戶空間庫中復制所有內容，不過，這就是為什么我們不能只繼承 BulkInferrer 並導入我們的自定義版本。）

我們必須確保'our_project:namespace.TableName'處的表具有與模型的 output 兼容的架構，但不必將該架構轉換為 JSON / AVRO。

理論上，我的團隊希望使用圍繞此構建的 TFX 提出拉取請求，但目前我們正在硬編碼幾個關鍵參數，並且沒有時間將其用於真正的公共/生產 state。

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"]
    )