Чередование файлов с подпроцессом. Открыть в tf.data

Question

У меня есть несколько файлов, которые я хотел бы использовать небольшими порциями до EOF с tf.data вместо использования tf.read_file один раз на файл (так как некоторые файлы намного больше других).

Я не знаю, как использовать конвейерные подпроцессы в качестве операции TensorFlow (tf.py_func как-нибудь?), А элемент набора данных из list_files известен только во время выполнения графа, поэтому следующее не работает:

def stream(path, bytesize=2048):
    args = f'my_program {path}'
    with subprocess.Popen(args, stdout=subprocess.PIPE) as pipe:
        while True:
            buffer = pipe.stdout.read(bytesize)
            yield np.frombuffer(buffer)
            if len(buffer) < bytesize:
                break


def map_func(path):
    generator = functools.partial(stream, path)
    dataset = tf.data.Dataset.from_generator(generator, tf.float32)
    return dataset


dataset = (
    tf.data.Dataset
    .list_files('data/*')
    .interleave(map_func, batch_size)
    .batch(batch_size)
)

Есть ли какой-то способ получить значение элемента набора данных в итерируемое значение, ожидаемое tf.data.Dataset.from_generator, или я поступаю неправильно?

Related: Может ли функция map, предоставленная для `tf.data.Dataset.from_generator (...)`, разрешить тензорный объект?

Answer 1

TensorFlow только что получил поддержку параметризованных генераторов в tf.data!

def map_func(path):
    dataset = tf.data.Dataset.from_generator(stream, tf.float32, args=(path,))
    return dataset

pip install tf-nightly или tf-nightly-gpu, чтобы попробовать вышесказанное.