tensorflow: tf.py_function in tf.data.Dataset pipeline doesn't work with TPUEstimator

Environment is Google Colab with TPU runtime.

Describe the current behavior I have an input pipeline that contains tf.py_function as one of its processing steps when training a model with TPUEstimator. When I run the code, I get the following error:

No registered 'EagerPyFunc' OpKernel for CPU devices compatible with node {{node EagerPyFunc}}
	.  Registered:  <no registered kernels>

	 [[EagerPyFunc]]
	 [[input_pipeline_task0/MakeIterator]]

According to the documentation (https://www.tensorflow.org/guide/using_tpu) “The input pipeline generated by your input_fn is run on CPU.” Running the same input code with a standard Estimator on CPU/GPU works just fine. Manually placing the Dataset + all processing steps on the CPU with tf.device('/cpu:0') also fails with the same error when training with TPUEstimator.

Describe the expected behavior I should be able to run Python code on the CPU as part of my input pipeline when training on TPUs.

Code to reproduce the issue https://colab.research.google.com/drive/15KuVkukIWdN753ffNjY6jVzT1gyhsBRQ

The notebook linked above is nearly identical to https://colab.research.google.com/github/tensorflow/tpu/blob/master/tools/colab/shakespeare_with_tpuestimator.ipynb. I’ve just added the following line to the Training Data code cell:

idx = tf.py_function(lambda x: x, [idx], tf.int32)

About this issue

Original URL
State: closed
Created 5 years ago
Reactions: 2
Comments: 19 (5 by maintainers)

Most upvoted comments

@jhseu with the TPUStrategy API I am still getting the same error as described in the initial issue. I am running on a Google Cloud v3-8, TF version 2.1.0-dev20200102, following these instructions. It works fine as is, but if I add the following two lines after line 86 to use py_functions:

inputs = tf.py_function(lambda x: x, [inputs], tf.int64)
targets = tf.py_function(lambda x: x, [targets], tf.int64

then I get this error:

tensorflow.python.framework.errors_impl.InvalidArgumentError: No OpKernel was registered to support Op 'EagerPyFunc' used by {{node EagerPyFunc}} with these attrs: [Tout=[DT_INT64], is_async=false, Tin=[DT_INT64], _xla_inferred_shapes=[<unknown>], token="pyfunc_0"]
Registered devices: [CPU, TPU, TPU_SYSTEM, XLA_CPU]
Registered kernels:
  <no registered kernels>

         [[EagerPyFunc]]

I will try to write my pipeline in terms of tf.data, but if you had any updates on using tf.py_function in this context, that would be great.

Alex-Fabbri on Jan 9, 2020

Yeah, in the new TPUStrategy API, we do already support this, and additionally it offers flexibility for the user to do step-at-a-time training and feed in whatever Python numpy arrays they want. It’s a bit new, so we’ll start documenting it soon.

jhseu on Jul 23, 2019

I have the same problem as above, any updates?

samkleeman1 on Mar 25, 2020

tf 2.3.0 does not work either.

fengyang0317 on Aug 14, 2020

Yeah, in the new TPUStrategy API, we do already support this, and additionally it offers flexibility for the user to do step-at-a-time training and feed in whatever Python numpy arrays they want. It’s a bit new, so we’ll start documenting it soon.

@jhseu Is there any updates? I tried to use tf.py_function inside dataset.map and still gets an error on TPU.

fengyang0317 on Jun 14, 2020