A small function for creating a windowed dataset for sequential networks using tensorflow 2.x

windowed_dataset.py

def windowed_dataset(series, window_size, batch_size, shuffle_buffer):
  """Function for creating a windowed dataset for sequence training"""
  dataset = tf.data.Dataset.from_tensor_slices(series)
  dataset = dataset.window(window_size + 1, shift=1, drop_remainder=True)
  dataset = dataset.flat_map(lambda window: window.batch(window_size + 1))
  dataset = dataset.shuffle(shuffle_buffer).map(lambda window: (window[:-1], window[-1]))
  dataset = dataset.batch(batch_size).prefetch(1)
  return dataset