使用 Tensorflow Transform 的 Apache Beam 作业（Python）被 Cloud Dataflow 杀死答案

【问题标题】：Apache Beam job (Python) using Tensorflow Transform is killed by Cloud Dataflow使用 Tensorflow Transform 的 Apache Beam 作业（Python）被 Cloud Dataflow 杀死
【发布时间】：2018-09-13 02:28:13
【问题描述】：

我正在尝试在 Dataflow 上运行基于 Tensorflow 转换的 Apache Beam 作业，但它被杀死了。有人经历过这种行为吗？这是 DirectRunner 的一个简单示例，在我的本地运行正常，但在 Dataflow 上失败（我正确更改了运行器）：

import os
import csv
import datetime
import numpy as np

import tensorflow as tf
import tensorflow_transform as tft

from apache_beam.io import textio
from apache_beam.io import tfrecordio

from tensorflow_transform.beam import impl as beam_impl
from tensorflow_transform.beam import tft_beam_io 
from tensorflow_transform.tf_metadata import dataset_metadata
from tensorflow_transform.tf_metadata import dataset_schema

import apache_beam as beam


NUMERIC_FEATURE_KEYS = ['feature_'+str(i) for i in range(2000)]


def _create_raw_metadata():
    column_schemas = {}
    for key in NUMERIC_FEATURE_KEYS:
        column_schemas[key] = dataset_schema.ColumnSchema(tf.float32, [], dataset_schema.FixedColumnRepresentation())

    raw_data_metadata = dataset_metadata.DatasetMetadata(dataset_schema.Schema(column_schemas))

    return raw_data_metadata


def preprocessing_fn(inputs):
    outputs={}

    for key in NUMERIC_FEATURE_KEYS:
        outputs[key] = tft.scale_to_0_1(inputs[key])

    return outputs


def main():

    output_dir = '/tmp/tmp-folder-{}'.format(datetime.datetime.now().strftime('%Y%m%d%H%M%S'))

    RUNNER = 'DirectRunner'

    with beam.Pipeline(RUNNER) as p:
        with beam_impl.Context(temp_dir=output_dir):

            raw_data_metadata = _create_raw_metadata()
            _ = (raw_data_metadata | 'WriteInputMetadata' >> tft_beam_io.WriteMetadata(os.path.join(output_dir, 'rawdata_metadata'), pipeline=p))

            m = numpy_dataset = np.random.rand(100,2000)*100
            raw_data = (p
                    | 'CreateTestDataset' >> beam.Create([dict(zip(NUMERIC_FEATURE_KEYS, m[i,:])) for i in range(m.shape[0])]))

            raw_dataset = (raw_data, raw_data_metadata)

            transform_fn = (raw_dataset | 'Analyze' >> beam_impl.AnalyzeDataset(preprocessing_fn))
            _ = (transform_fn | 'WriteTransformFn' >> tft_beam_io.WriteTransformFn(output_dir))

            (transformed_data, transformed_metadata) = ((raw_dataset, transform_fn) | 'Transform' >> beam_impl.TransformDataset())

            transformed_data_coder = tft.coders.ExampleProtoCoder(transformed_metadata.schema)
            _ = transformed_data | 'WriteTrainData' >> tfrecordio.WriteToTFRecord(os.path.join(output_dir, 'train'), file_name_suffix='.gz', coder=transformed_data_coder)

if __name__ == '__main__':
  main()

另外，我的生产代码（未显示）失败并显示以下消息：The job graph is too large. Please try again with a smaller job graph, or split your job into two or more smaller jobs.

有什么提示吗？

【问题讨论】：

Dataflow 只支持小于（我相信是 10MB 或 100MB）的管道图。此错误意味着您的图表太大。让我看看我们是否能找到解决方法...
我在使用模板方法时遇到了同样的“太大”错误。限制为 10MB。改用传统方法克服了那个特定的错误，然后又遇到了另一个限制。这次是 20MB 的限制。这是由于其他一些 .json 文件变得太大。

标签： tensorflow google-cloud-dataflow apache-beam tensorflow-transform

【解决方案1】：

此处记录了对管道描述大小的限制： https://cloud.google.com/dataflow/quotas#limits

有一种方法可以解决这个问题，我们可以先将它们堆叠在一起，然后使用“elementwise=True”将它们传递到 tft.scale_to_0_1 来融合它们，而不是为进入 tft.scale_to_0_1 的每个张量创建阶段。

结果将是相同的，因为最小值和最大值是按“列”而不是整个张量计算的。

这看起来像这样：

stacked = tf.stack([inputs[key] for key in NUMERIC_FEATURE_KEYS], axis=1)
scaled_stacked = tft.scale_to_0_1(stacked, elementwise=True)
for key, tensor in zip(NUMERIC_FEATURE_KEYS, tf.unstack(scaled_stacked, axis=1)):
  outputs[key] = tensor

【讨论】：