Tensorflow 数据集交错 from_generator 抛出 InvalidArgumentError

我有一个生成器，我正在尝试交错：

def hello(i):
  for j in tf.range(i):
    yield j

ds = tf.data.Dataset.range(10).interleave(
       lambda ind: tf.data.Dataset.from_generator(lambda: hello(ind), output_types=(tf.int32,)))

for x in ds.take(1):
  print(x)

但我收到这个错误：

TypeError: An op outside of the function building code is being passed
a "Graph" tensor. It is possible to have Graph tensors
leak out of the function building context by including a
tf.init_scope in your function building code.
For example, the following function will fail:
  @tf.function
  def has_init_scope():
    my_constant = tf.constant(1.)
    with tf.init_scope():
      added = my_constant * 2
The graph tensor has name: args_0:0


     [[{{node PyFunc}}]]

张量流版本：2.3.2

问题在于您构建生成器函数的方式。而不是使用lambda，你应该使用args关键字参数指定传递给生成器函数的参数。

ds = tf.data.Dataset.range(10).interleave(
    lambda ind: tf.data.Dataset.from_generator(
        hello, args=(ind,), output_types=tf.int32
    )
)

对于 TF2.4，请注意您应该使用output_signature代替output_types，因为后者已被弃用。 （在这种情况下output_signature=tf.TensorSpec(shape=(), dtype=tf.int32,)).