Takes a tf.Dataset and creates a numpy stream of ready batches.

def dataset_to_stream(dataset, input_name, num_chunks=0, append_targets=False):
  """Takes a tf.Dataset and creates a numpy stream of ready batches."""
  for example in tfds.as_numpy(dataset):
    inp, out = example[0][input_name], example[1]
    if len(out.shape) > 1 and out.shape[-1] == 1:
      out = np.squeeze(out, axis=-1)
    if num_chunks > 0:
      inp = np.split(inp, num_chunks, axis=1)
      out = np.split(out, num_chunks, axis=1)
    if append_targets:
      inp = (inp, out)
    yield inp, out