python - 在带有 DataFlow 的 preprocessing_fn 中使用 tf.py_function 会导致 ValueError: callback pyfunc_16 is not found (Github 中的代码)

Question

我正在尝试在 Kubeflow Pipelines 上为 TFX运行这个Keras 组件DataflowRunner笔记本，但使用and 并稍作改动。在preprocessing_fn我替换tft.scale_to_z_score(_fill_in_missing(inputs[key]))为以下内容：

from sklearn.preprocessing import scale

def sk_scale_to_z_score(X):
  scaled = scale(X)
  return scaled

@tf.function
def tf_scale_to_z_score(X):
  X_shape = X.shape
  [X_scaled, ] = tf.py_function(sk_scale_to_z_score, [X], [tf.float32])
  X_scaled.set_shape(X_shape)
  return X_scaled

所以最终的代码片段如下所示：

outputs[_transformed_name(key)] = tf_scale_to_z_score(_fill_in_missing(inputs[key]))

但是，此管道在引发以下错误的Transform步骤中停止：Dataflow

 ValueError: callback pyfunc_16 is not found

因此，我的问题是：如果我正在使用，tf.py_function那么我可以使用Dataflow吗？如果是，那么如何？

编辑：代码可以在这个github repo中找到。

标签： pythontensorflowgoogle-cloud-dataflowapache-beamtfx