def dataset_to_stream(dataset, input_name):
"""Takes a tf.Dataset and creates a numpy stream of ready batches."""
# All input-pipeline processing should be on CPU.
for example in fastmath.dataset_as_numpy(dataset):
features = example[0]
inp, out = features[input_name], example[1]
mask = features['mask'] if 'mask' in features else None
# Some accelerators don't handle uint8 well, cast to int.
if isinstance(inp, np.uint8):
inp = inp.astype(np.int32)
if isinstance(out, np.uint8):
out = out.astype(np.int32)
yield (inp, out) if mask is None else (inp, out, mask)
def gen(generator=None):
del generator
for example in fastmath.dataset_as_numpy(dataset):
yield example
def gen(unused_arg):
for example in fastmath.dataset_as_numpy(dataset):
yield example
