tf.data.Dataset.from_generator()字符串转字节问题求助

问题复现 #

当使用tf.data.Dataset.from_generator()时，传入的字符串参数会被自动转为字节串（bytes类型），导致生成器函数中原本预期的字符串参数变成字节对象，进而引发逻辑错误。复现代码及输出如下：

import tensorflow as tf
from tensorflow.data import Dataset

def some_generator(date1, date2):
    print(date1, date2)
    yield [1, 2], [3]

feats, label = next(some_generator('2010-01-01', '2017-12-31'))
signature = (tf.type_spec_from_value(tf.convert_to_tensor(feats)),
             tf.type_spec_from_value(tf.convert_to_tensor(label)))
ds = Dataset.from_generator(some_generator, args=('2010-01-01', '2017-12-31'), output_signature=signature)

for feats, label in ds.take(1):
    print(feats, label)

输出：

2010-01-01 2017-12-31
b'2010-01-01' b'2017-12-31'
tf.Tensor([1 2], shape=(2,), dtype=int32) tf.Tensor([3], shape=(1,), dtype=int32)

解决方案 #

方案1：在生成器内部将字节串转回字符串

直接在生成器开头判断参数类型，若为字节串则解码为字符串：

def some_generator(date1, date2):
    # 处理字节串转字符串
    if isinstance(date1, bytes):
        date1 = date1.decode('utf-8')
    if isinstance(date2, bytes):
        date2 = date2.decode('utf-8')
    print(date1, date2)
    yield [1, 2], [3]

这种方法兼容性强，不管参数是直接传入字符串还是被tf.data转为字节串，都能正常处理。

方案2：用lambda包装生成器，提前绑定参数

通过lambda函数预先固定生成器的参数，避免tf.data对参数进行类型转换：

ds = Dataset.from_generator(
    lambda: some_generator('2010-01-01', '2017-12-31'),
    output_signature=signature
)

此时from_generator调用的是无参lambda，lambda内部直接调用带字符串参数的生成器，tf.data不会干预参数的类型。

原因说明 #

tf.data在处理from_generator的args参数时，会将Python原生类型转为对应的TensorFlow张量类型：字符串会被转为tf.string张量，对应到Python中就是bytes对象；而整数等类型的张量与Python原生类型兼容度更高，因此不会出现类型转换问题。

内容的提问来源于stack exchange，提问作者mac13k