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 math.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 dataset_to_stream(dataset, input_name):
"""Takes a tf.Dataset and creates a numpy stream of ready batches."""
for example in math.dataset_as_numpy(dataset):
features = example[0]
inp, out = features[input_name], example[1]
mask = features['mask'] if 'mask' in features else None
# All input-pipeline processing should be on CPU.
with tf.device('cpu:0'):
# Some accelerators don't handle uint8 well, cast to int.
if isinstance(inp, jnp.uint8):
inp = inp.astype(jnp.int32)
if isinstance(out, jnp.uint8):
out = out.astype(jnp.int32)
if len(out.shape) > 1 and out.shape[-1] == 1:
out = jnp.squeeze(out, axis=-1)
yield (inp, out) if mask is None else (inp, out, mask)