# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     name: python3
# ---

# + [markdown] id="view-in-github" colab_type="text"
# <a href="https://colab.research.google.com/github/adholmgren/tensorflow_playground/blob/master/tf_dataset.ipynb" target="_parent"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/></a>

# + [markdown] id="RIhbCtak8CYS" colab_type="text"
# Copyright 2020, <NAME>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

# + [markdown] id="skZshN7Y85YR" colab_type="text"
# # A practical guide to the Dataset class in tensorflow
# The tensorflow Dataset class can streamline your machine learning application, but I've personally found many times when the documentation could benefit from some more examples or some more exposition. This notebook will hopefully serve as a guide, or at the very least a supplement, to the tensorflow documentation. An intermediate level of Python knowledge is assumed but not necessarily required, some concepts are best communicated with intermediate Python knowledge. Basic Python is required.

# + id="csllqN8Jpjar" colab_type="code" colab={}
try:
  # # %tensorflow_version only exists in Colab.
  # %tensorflow_version 2.x
except Exception:
  # !pip install --quiet tensorflow-gpu>=2.0.0

# + id="6D8PWnkzlE7u" colab_type="code" colab={}
import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
import tensorflow.keras as K
# tf.enable_eager_execution()

# + [markdown] id="cW6XGcP6nSWy" colab_type="text"
# # Build dataset from arrays

# + [markdown] id="sD-rO2aUGZwS" colab_type="text"
# If the data comes in as arrays it's pretty simple to get the data into a tensorflow dataset using the `from_tensor_slices` method.

# + id="xbAEiXD_lz12" colab_type="code" colab={}
arr = np.array([7, 2, 1, 6, 3, 5, 9])
dataset = tf.data.Dataset.from_tensor_slices(arr)

# + [markdown] id="uPMpOJcg4V1a" colab_type="text"
# The tensor slices assumes the first dimension of the array is the dimension of new instances. For example, in the last code block it consumed the (7,) array and generated 7 instances of () shaped Tensors. As another example, the following code will consume 2 instances of 2x2 arrays.

# + id="8C1H8HnD4WWb" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 126} outputId="3e2b51ae-f20a-4742-f428-c939d9a49a6e"
dataset_2D = tf.data.Dataset.from_tensor_slices(np.reshape(np.arange(2**3), (2, 2, 2)))
for elem in dataset_2D:
    print(elem)

# + [markdown] id="L8FUm9JnEAbp" colab_type="text"
# In comparison, this code will give just one instance of a 2x2x2 array.

# + id="I0izNSFjD5d9" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 126} outputId="16ac13fd-ddfc-4911-8804-f8517151ff84"
dataset_3D = tf.data.Dataset.from_tensor_slices(np.reshape(np.arange(2**3), (1, 2, 2, 2)))
for elem in dataset_3D:
    print(elem)

# + [markdown] id="jg8kpRN0EGSh" colab_type="text"
# **Test yourself**: build a dataset with 3 instances of a 3x3 array. Build 1 instance of a 3x10 array.

# + id="JUl7LP1wQifP" colab_type="code" colab={}
# 3 instances of 3x3 array
# code here

# 1 instance of a 3x10 array
# code here

# + [markdown] id="2dicUb1SQzP4" colab_type="text"
# Personally, I feel this next method is a bit redundant (I haven't looked into the source code, there could some slightly different optimization), but there's also a `from_tensor` method that creates a single instance. I think the context that makes the most sense for this method is testing at inference.

# + id="wqIVSGELQyUX" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 126} outputId="295d971c-9f27-431d-ac68-d5d75e182465"
dataset_3D = tf.data.Dataset.from_tensors(np.reshape(np.arange(2**3), (2, 2, 2)))
for elem in dataset_3D:
    print(elem)

# + [markdown] id="C0LdGDFtUwaB" colab_type="text"
# You can see that this is the same thing as using  
# ```Python
# dataset_3D = tf.data.Dataset.from_tensor_slices(np.reshape(np.arange(2**3), (1, 2, 2, 2)))
# ```  
# Basically, if your array is a single instance this is a way for the Dataset class to consume the single instance. 

# + [markdown] id="bPeOb7imIPUt" colab_type="text"
# # Using methods in the dataset class

# + [markdown] id="kqAg3cZPRfZ8" colab_type="text"
# ### Dataset as an iterator

# + [markdown] id="VcISuxtRu19b" colab_type="text"
# Just like most Python objects, the dataset class has an iterator. In fact, the fundamental point of the dataset class is that its very essence is to be an iterator. All the methods within the dataset class either instantiate the iterator (e.g. `from_tensor_slices`), modify the iterator (e.g. `map`), or control the iterator (e.g. `batch`). The Dataset class gives users flexibility to control the memory and processing in flowing data inputs to Tensorflow's machine learning models. Most people can, and probably want to, stop there -- the dataset is fundamentally an iterator similar to Python's range, or numpy arrays, or a thousand other Python objects that iterate.
#
# For anyone interested in peaking a bit more into the nitty gritties read on. If you're not familiar with Python iterators and how they're built, here's a pretty good guide ([link](https://www.ics.uci.edu/~pattis/ICS-33/lectures/iterators.txt)). To really find out what's happening with the dataset iterator you have to go to the source code. The first thing that the source code lets you know is that it implements the Python iterator protocol and therefore can only be used in eager mode. This comment actually has the potential to cause confusion because the Dataset class can be used with a static graph, it will just function differently. An important consequence of eager vs. static execution, is that in static graph the entirety of arrays are placed into the static graph as Variables (potentially taking up a lot of memory and running into byte limits in graph serialization).
# ```Python
# def __iter__(self):
#     """Creates an `Iterator` for enumerating the elements of this dataset.
#     The returned iterator implements the Python iterator protocol and therefore
#     can only be used in eager mode.
#     Returns:
#       An `Iterator` over the elements of this dataset.
#     Raises:
#       RuntimeError: If not inside of tf.function and not executing eagerly.
#     """
#     if (context.executing_eagerly()
#         or ops.get_default_graph()._building_function):  # pylint: disable=protected-access
#       return iterator_ops.OwnedIterator(self)
#     else:
#       raise RuntimeError("__iter__() is only supported inside of tf.function "
#                          "or when eager execution is enabled.")
# ```
# We then see that it uses the tensorflow iterator_ops for its iterator, so we can go to that source code. That source code points us to the [iterator_ops](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/data/ops/iterator_ops.py). Of which, the main elements are `_create_iterator` and `_next_internal`. These are effective the `__iter__` and `__next__` methods for the class. Those functions themselves have some more nitty gritties, but suffice to say that you can trace it far enough to reassure yourself that the iterator is moving through nests of the Tensor class. The gist of `_create_iterator` is that it has some pretty thorough robustness so you don't break it, but ultimately it is assigning pieces that fundamentally want to work with tensors. The `_next_internal` checks to see whether the backend is in eager_execution or static graph execution. In eager execution, where you'll most easily see the results, the next method is primarily concerned with returning tensor elements. There's a lot more tracking and integrating the core fundamentals than described thus far, but hopefully this at least provides a shallow insight into the class.
# ```Python
# def _create_iterator(self, dataset):
#     # pylint: disable=protected-access
#     dataset = dataset._apply_options()
#
#     # Store dataset reference to ensure that dataset is alive when this iterator
#     # is being used. For example, `tf.data.Dataset.from_generator` registers
#     # a few py_funcs that are needed in `self._next_internal`.  If the dataset
#     # is deleted, this iterator crashes on `self.__next__(...)` call.
#     self._dataset = dataset
#
#     ds_variant = dataset._variant_tensor
#     self._element_spec = dataset.element_spec
#     self._flat_output_types = structure.get_flat_tensor_types(
#         self._element_spec)
#     self._flat_output_shapes = structure.get_flat_tensor_shapes(
#         self._element_spec)
#     with ops.colocate_with(ds_variant):
#       self._iterator_resource, self._deleter = (
#           gen_dataset_ops.anonymous_iterator_v2(
#               output_types=self._flat_output_types,
#               output_shapes=self._flat_output_shapes))
#       gen_dataset_ops.make_iterator(ds_variant, self._iterator_resource)
#       # Delete the resource when this object is deleted
#       self._resource_deleter = IteratorResourceDeleter(
#           handle=self._iterator_resource,
#           device=self._device,
#           deleter=self._deleter)
#
# def _next_internal(self):
#     """Returns a nested structure of `tf.Tensor`s containing the next element.
#     """
#     if not context.executing_eagerly():
#       with ops.device(self._device):
#         ret = gen_dataset_ops.iterator_get_next(
#             self._iterator_resource,
#             output_types=self._flat_output_types,
#             output_shapes=self._flat_output_shapes)
#       return structure.from_compatible_tensor_list(self._element_spec, ret)
#
#     # This runs in sync mode as iterators use an error status to communicate
#     # that there is no more data to iterate over.
#     # TODO(b/77291417): Fix
#     with context.execution_mode(context.SYNC):
#       with ops.device(self._device):
#         # TODO(ashankar): Consider removing this ops.device() contextmanager
#         # and instead mimic ops placement in graphs: Operations on resource
#         # handles execute on the same device as where the resource is placed.
#         ret = gen_dataset_ops.iterator_get_next(
#             self._iterator_resource,
#             output_types=self._flat_output_types,
#             output_shapes=self._flat_output_shapes)
#
#       try:
#         # Fast path for the case `self._structure` is not a nested structure.
#         return self._element_spec._from_compatible_tensor_list(ret)  # pylint: disable=protected-access
#       except AttributeError:
#         return structure.from_compatible_tensor_list(self._element_spec, ret)
# ```

# + [markdown] id="oN_ZeO66Zz4F" colab_type="text"
# Here's the most common way Dataset iterator is used

# + id="wH3Ju6JQ8i96" colab_type="code" colab={}
# create dataset
arr = np.array([7, 2, 1, 6, 3, 5, 9])
dataset = tf.data.Dataset.from_tensor_slices(arr)

# + id="X22Omo5ktP1E" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 144} outputId="6bb2d589-36fe-44fa-c75c-0350b38bf805"
# loop through dataset elements
for elem in dataset:
    print(elem)

# + [markdown] id="f-M2JnN9aET1" colab_type="text"
# If you only wanted a certain number of elements, you could even set conditional break in there.

# + id="zpyJP4GLaS-s" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 72} outputId="07c2529f-d45e-49b0-ec87-ac94fca490a6"
count = 0
n_elem = 3
for elem in dataset:
    print(elem)
    count += 1
    if count >= n_elem:
        break

# + [markdown] id="VvXu5tSW-diS" colab_type="text"
# Another way to see just the first few elements is to work with the bare iterator.

# + id="UQ3-rVgg-ZQF" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 72} outputId="a7c51cff-9b5f-4144-879b-741cb9701e5e"
ds_iter = iter(dataset)
print(f'first element: {next(ds_iter)}')
print(f'second element: {next(ds_iter)}')
print(f'third element: {next(ds_iter)}')
del ds_iter

# + [markdown] id="5xDKViKc9aH5" colab_type="text"
# If you want to know what the for loop is really doing, we could do the functionally same operation written as below. The key is that the for loop goes until it hits a StopIteration exception that is built into the class's `next` method.

# + id="GwkwEnyU9Y_6" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 144} outputId="0c2f54db-912b-4958-85ec-eaad0b41cf4e"
ds_iter = iter(dataset)
try:
    while True:
        elem = next(ds_iter)
        print(elem)
except StopIteration:
        pass
finally:
        del ds_iter

# + [markdown] id="Ll95DKJQZt8L" colab_type="text"
# Not as concise, is it?

# + [markdown] id="aRvsKP79D1vA" colab_type="text"
# ## Take

# + [markdown] id="C7jElLnusN4g" colab_type="text"
# Take, grab, give, this could have been called a lot of things but, alas, it is called `take`. This method takes a count (integer) and creates a new dataset that will iterate through tensors from at most count elements. The statement to reiterate is take builds a new dataset, it does not extract elements. Another tip worth noting is that you can overcount the take count and when the dataset iterates it will just go to max number of elements in the dataset. Similarly, any negative number simply takes the entire dataset.

# + id="tmh7HNljsMZA" colab_type="code" colab={}
arr = np.array([7, 2, 1, 6, 3, 5, 9])
dataset = tf.data.Dataset.from_tensor_slices(arr)

# + id="4nI1q6O0trqB" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 108} outputId="ea3fdc73-ff7c-40f0-edd3-87332a9b38da"
ds_take = dataset.take(3)
print(list(ds_take.as_numpy_iterator()), "\n")
ds_take_more = dataset.take(1000)
print("take over count: ", list(ds_take_more.as_numpy_iterator()))
ds_take_neg = dataset.take(-8)
print("take negative count (take all): ", list(ds_take_neg.as_numpy_iterator()))
ds_take_none = dataset.take(0)
print("take zero count (none): ", list(ds_take_none.as_numpy_iterator()))


# + [markdown] id="YL5qQ_UivWCv" colab_type="text"
# A decent analogy for `take` would be a wrapper on a generator with a stop condition that makes the first generator stop earlier if the count is reached and keep going if not.  Since the [true class](https://github.com/tensorflow/tensorflow/blob/e26286914b6d3cf7ed0c9a47f50c07391c2174a6/tensorflow/core/kernels/data/take_dataset_op.cc) is defined in C++, this analogy is merely illustrative.

# + id="COwMHtrQvUev" colab_type="code" colab={}
# Similar concept to Dataset.take
def foo_generator():  # base iterator like Dataset
    for i in range(10):
        yield i 

def foo_take(count):
    counter = 0
    pity_the_foo = foo_generator()
    run_all = count < 0
    while True:
        try:
            if (counter >= count) and not run_all:
                raise StopIteration
            yield next(pity_the_foo)
            counter += 1
        except StopIteration:
            break


# + id="Karmjjc9wsyu" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 90} outputId="e6cc36c3-4dd0-40ff-b85e-2750e3c007cf"
print(list(foo_take(5)))
print(list(foo_take(1000)))
print(list(foo_take(-1)))
print(list(foo_take(0)))

# + [markdown] id="ZFJMJMGzQMnu" colab_type="text"
# ## Map dataset

# + id="SNkftOO-8kyR" colab_type="code" colab={}
arr = np.array([7, 2, 1, 6, 3, 5, 9])
dataset = tf.data.Dataset.from_tensor_slices(arr)


# + [markdown] id="g72VAcA9QWtV" colab_type="text"
# You can use map to apply a function to the dataset. For example, if you want to add 1 to to every element of the dataset you could do

# + id="jZ1iIrvVmCR6" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 162} outputId="7b449f2f-3222-4ff6-dd06-2fe391a94d70"
def add_one(x):  # note: easily done anonymously lambda x: x + 1
    return x + 1 

dp1 = dataset.map(add_one)
print("dataset ", "mapped dataset")
for ds_elem, elem in zip(dataset, dp1):
    print(ds_elem.numpy(), "\t", elem.numpy())


# + [markdown] id="vyV1gQn_Qng0" colab_type="text"
# As another example of a mapping function, we can make a function that one-hot encodes values corresponding to a class index.

# + id="MR0jPfHLszlV" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 144} outputId="7cd79766-2d87-4d3f-ab81-d2b4fccbce08"
def one_hot_encode(x, max_index):
    tensor_zeros = tf.zeros(max_index, dtype=tf.int32)
    tensor_ones = tf.ones(max_index, dtype=tf.int32)
    tensor_range = tf.range(max_index, dtype=tf.int32)
    # dataset defaulted to int64 type, hence the casting, types may vary by context
    one_hot_array = tf.where(tensor_range == tf.cast(x, tf.int32), tensor_ones, tensor_zeros)
    return one_hot_array

ds_one_hot = dataset.map(lambda x: one_hot_encode(x, tf.constant(10, dtype=tf.int32)))
for elem in ds_one_hot:
    print(elem)

# + [markdown] id="gWjBfpr4pW88" colab_type="text"
# Alternatively, rather than manually making a one hot encoding function you could just use tensorflows built-in one_hot function.

# + id="BHXuwuGCjcyh" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 144} outputId="bc43e8ee-5469-4f61-d047-c5ba628e719e"
ds_one_hot = dataset.map(lambda x: tf.one_hot(x, tf.constant(10, dtype=tf.int32), dtype=tf.int32))
for elem, one_hot_elem in zip(dataset, ds_one_hot):
    print(f'category index {elem.numpy()}, one_hot_encoding {one_hot_elem.numpy()}')


# + [markdown] id="h_Wyy0xdAUgN" colab_type="text"
# To really do the one-hot encoding to death, note that the tensors could also return as a sparse type.

# + id="wnZLPDlNAJui" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 291} outputId="40a7fe62-da13-4409-e411-596065a95e67"
def one_hot_encode(x, max_index):
    return tf.sparse.SparseTensor(indices=[[x]], values=[tf.cast(1, dtype=tf.int64)], dense_shape=[max_index])

ds_one_hot = dataset.map(lambda x: one_hot_encode(x, tf.constant(10, dtype=tf.int64)))
for elem in ds_one_hot:
    print(elem)
    print(tf.sparse.to_dense(elem))


# + [markdown] id="5tDCdbYpDBVV" colab_type="text"
# ## Filter dataset

# + [markdown] id="pQULfD0eIKwy" colab_type="text"
# Both map and filter applies a function to each element in a dataset. Wheras map uses a function to alter the dataset elements in some way, filter works on a conditional return to delete dataset elements that fail the conditional test. The filter method is good for deleting bad dataset inputs.  
#
# Let's imagine a dataset of images that vary in length and width. Most pretrained networks are trained for images of sizes 299x299, 256x256, or 224x224. So if images are coming in the general size of 300x300, some more (say 350x350) and some less (say 250x250) then for the most part I can safely crop down to the right size. However, if the data augmentation involves rotation or other affine transforms then extra support around the input size will be desired, and as such some images may need to get thrown out. The following example goes through such a case.

# + id="fg7gTTw7M_cZ" colab_type="code" colab={}
# generate random array of random size
def gen_series(n_samps):
  i = 0
  while True:
    rand_shape = np.random.randint(2, 5)
    # note need to explicitly output shape because a symbolic tensor does not have a shape 
    yield (rand_shape, np.random.random(size=(rand_shape, rand_shape)))
    i += 1
    if i > n_samps:
        break



# + [markdown] id="faqKiLJW83q5" colab_type="text"
# Now set the random seed for reproducibility and make 5 random numpy arrays.

# + id="fMjtD0N1JvIZ" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 380} outputId="0aee30b7-d45a-4315-a663-220bebe9c2be"
np.random.seed(314)
for elem in gen_series(5):
    print(elem)

# + [markdown] id="4mlISeCvKvAR" colab_type="text"
# Turn the arrays into a dataset

# + id="itYcFMCtKuGT" colab_type="code" colab={}
np.random.seed(314)
dataset_shapes = tf.data.Dataset.from_generator(gen_series,
                                                output_types=(tf.int32, tf.float32),
                                                args = [5]
                                                )


# + [markdown] id="nA-sGx0o9Dsm" colab_type="text"
# Filter out any arrays that have a shape less than 3 (or conversely keep arrays larger than 2)

# + id="_6E99bIQgg7K" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 344} outputId="ffb75228-05c9-4c11-c052-47391df9337a"
def filter_shape(shape, arr):
    return shape > 2

np.random.seed(314)
ds_filter = dataset_shapes.filter(filter_shape)
list(ds_filter.as_numpy_iterator())

# + [markdown] id="oAtrDF3c0PUm" colab_type="text"
# Another case where a filter could be helpful is if the data hits some value then you know it's corrupt (maybe a CRC check, or just physical constraints). For example, it could be that you know its impossible for a speed-o-meter to register a speed less than 0.

# + id="YHLrrOxp0oiu" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 199} outputId="bbfd26f5-b6a3-434a-b183-9110ce790c20"
np.random.seed(272)
dataset_threshold = tf.data.Dataset.from_tensor_slices(np.random.rand(10)).map(lambda x: x - 0.3)
list(dataset_threshold.as_numpy_iterator())

# + [markdown] id="hEvqimxT-wsm" colab_type="text"
# Keep any values that are above 0

# + id="tGm6FzAC-hMV" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 108} outputId="f21b0266-df51-4806-bd07-d009faab57e9"
list(dataset_threshold.filter(lambda x: x > 0).as_numpy_iterator())

# + [markdown] id="oo6shlT1-9k_" colab_type="text"
# Try it out for yourself. Make an array with normally distributed data and filter out an events more than 2*sigma.

# + [markdown] id="2CiKDE7gWolw" colab_type="text"
# ## Interleave

# + [markdown] id="XIpjB0EPW7TJ" colab_type="text"
# Building from the concepts of `map`, `interleave` is `map` operation. What is different, though, is that `interleave` is a map of a dataset. Wheras `map` operates on tensor elements and typically involves processing those elements, the `interleave` method operates on a dataset itself and typically involves creating more samples. In this sense interleave often builds on a map.

# + [markdown] id="nQIfQGLSa07g" colab_type="text"
# Personally, I found this function was really difficult to wrap my head around.
#
# Interleave has two main controls, `cycle_length` and `block_length`.   
#   * `cycle_length` 
#   * `block_length`
#
# Cycle length is how many times to cycle through the datasets before starting to interleave, and block length is the length of the interleave.

# + id="Mywyost_LyDZ" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 562} outputId="a6bf980f-642d-41d6-db7c-ab23772ea28d"
dataset = tf.data.Dataset.range(1, 6)  # ==> [ 1, 2, 3, 4, 5 ]
# NOTE: New lines indicate "block" boundaries.
dataset = dataset.interleave(
    lambda x: tf.data.Dataset.from_tensors(x).repeat(6),
    cycle_length=3, block_length=4)
list(dataset.as_numpy_iterator())


# + id="lA9skUBgXPHT" colab_type="code" colab={}
base_arr = [['a'], ['b']]
num_array = [np.arange(4), np.arange(6, 10)]

# + id="yNRWuqeIez2m" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 35} outputId="f38e1d6c-9dab-43b6-9b7a-b531b087a9b8"
transform = lambda x: -x
dataset = tf.data.Dataset.from_tensor_slices(base_arr)
print(list(dataset.as_numpy_iterator()))

# + id="PWTJogF5Y8QQ" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 794} outputId="517c7f06-51e8-4eb3-b31e-41a66ae0a635"
interleaved_set = dataset.interleave(lambda x: tf.data.Dataset.from_tensor_slices(num_array),
                                     cycle_length=1,
                                     block_length=1,
                                     )
list(interleaved_set.as_numpy_iterator())

# + [markdown] id="TkPLZBDlW7zQ" colab_type="text"
# ## Shuffle

# + [markdown] id="chl095YwxMJ9" colab_type="text"
# Shuffle is pretty much what you imagine, it takes your dataset and will spit out a random permutation of your dataset. The main nuance to the shuffle function is the buffer size argument. That being said, the buffer size argument is itself pretty nuanced.

# + [markdown] id="H7EaryvYzvYG" colab_type="text"
# Take the following example: there's a dataset with numbers one through 9.  
# With a buffer size of 2, the first shuffle picks from  
# [0, 1]  
# let's say that it picks 1, now the buffer has  
# [0, 2]  
# let's say that it picks 2 this time around, now the buffer has  
# [0, 3]  
# and the shuffled array, to date, has  
# [1, 2, 3]  
# elements. As such, the placement of 0 in later spots of the array follows a geometric distribution. On the other end, 9 will only ever be in the 8th or 9th index (zero-indexing), or more generally for a shuffle size of 2 a value in the nth index can only appear in the n-1 index or the nth index. Even more generally, for shuffle size of s, a value in the nth index (zero-indexing) can appear in the max(0, n-s+1) index, e.g. with a shuffle size of 10 the 10th element or 9 index can appear at position 9-10+1=0 index.

# + id="pwiqEFWNqkDl" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 399} outputId="674b418d-b4b2-436e-89e2-b66f89c1005a"
ds = tf.data.Dataset.from_tensor_slices(np.arange(10))
print([x.numpy() for x in ds])
ds_shuffle = ds.shuffle(2, reshuffle_each_iteration=True)
for _ in range(20):
    print([x.numpy() for x in ds_shuffle])

# + [markdown] id="dvHPeW3O-_wV" colab_type="text"
# Another example of shuffle, but with a tuple input

# + id="uPaA5q7b_EMy" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 199} outputId="b8d0da82-f80e-46fa-d8c7-510f9be95e53"
rand_ds = tf.data.Dataset.from_tensor_slices((np.random.randn(10, 5), np.arange(10)))
for arr, label in rand_ds:
  print(arr.numpy(), '\t', label.numpy())

# + id="L630KzEw_Z8Q" colab_type="code" colab={}
for i in range(10):
  ds_iter = iter(rand_ds.shuffle(2, reshuffle_each_iteration=True))
  arr, label = next(ds_iter)
  print(arr.numpy(), '\t', label.numpy())
  arr, label = next(ds_iter)
  print(arr.numpy(), '\t', label.numpy())
  print('\n')

# + [markdown] id="pa4v-SjpnYIr" colab_type="text"
# # Datasets with x, y

# + id="T-XFtZiyncz1" colab_type="code" colab={}
train, test = tf.keras.datasets.fashion_mnist.load_data()


# + id="Ka7zRUBW5C0m" colab_type="code" colab={}
images, labels = train
images = images/255.0
labels = labels.astype(np.int32)

# + id="abrtBhLU5gbX" colab_type="code" colab={}
mnist_ds = dataset.shuffle(5000).batch(32)

# + id="quxNWfXan6pq" colab_type="code" colab={}
mnist_ds

# + id="u_yww7s26L4k" colab_type="code" colab={}
images_reshaped = images.reshape((60000, -1))

# + id="_gvtS2cD_gPj" colab_type="code" colab={}
type(images_reshaped)

# + id="jL8j1V3F6Xaj" colab_type="code" colab={}
images_reshaped.shape

# + id="K-LzTcvY6ZXT" colab_type="code" colab={}
ds_linear = tf.data.Dataset.from_tensor_slices((images_reshaped))

# + id="fAjf3Rtq6guU" colab_type="code" colab={}
mnist_ds_linear = ds_linear.shuffle(10).batch(1)

# + id="538UZyRu6oqj" colab_type="code" colab={}
mnist_ds_linear

# + id="3oHp9DSh6q3T" colab_type="code" colab={}
for num_now in ds_linear.take(2):
  plt.figure(); plt.imshow(np.reshape(num_now, (28, 28)), cmap='gray')

# + id="D1uwEbjl7QAo" colab_type="code" colab={}
for num_now in ds_linear.shuffle(5000).take(2):
  plt.figure(); plt.imshow(np.reshape(num_now, (28, 28)), cmap='gray')

# + id="Bvx04TsR5ZxX" colab_type="code" colab={}
fmnist_train_ds = tf.data.Dataset.from_tensor_slices((images, labels))
fmnist_train_ds = fmnist_train_ds.shuffle(5000).batch(32)

# + id="6e4SxXpmOgzw" colab_type="code" colab={}
model = tf.keras.Sequential([
  tf.keras.layers.Flatten(),
  tf.keras.layers.Dense(10)
])

# + id="vV1BmXWBO_Ti" colab_type="code" colab={}
model.compile(optimizer='adam',
              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True), 
              metrics=['accuracy'])


# + id="4mbdR45XPARV" colab_type="code" colab={}
model.fit(fmnist_train_ds, epochs=10)

# + id="jSAORFjzPSYr" colab_type="code" colab={}
result = model.predict(fmnist_train_ds, steps = 1)

# + id="Fljyfl1BPU-D" colab_type="code" colab={}
test_ds = tf.data.Dataset.from_tensor_slices((images, labels))

# + id="B9dcQwHlRX3B" colab_type="code" colab={}
x, y = list(zip(*list(test_ds.batch(5).take(1).as_numpy_iterator())))

# + id="ff7DBvU_RZsJ" colab_type="code" colab={}
y

# + id="Nsv_NPJzQBZc" colab_type="code" colab={}
x, y = list(zip(*list(test_ds.batch(5).take(1).as_numpy_iterator())))
y_predict_all = model.predict(test_ds.batch(5), steps=1)
y_top = np.argmax(y_predict_all, axis=1)
print(y, y_top)

# + id="MaCEj-jIQS01" colab_type="code" colab={}
for image_now in x[0]:
    plt.figure()
    plt.imshow(image_now.squeeze(), cmap='gray')

# + id="CUImNHuL2sGk" colab_type="code" colab={}
x_lin = tf.range(0, 100, dtype=tf.float32)
y_lin = tf.range(0, 100, dtype=tf.float32)

# + id="MUMJKSAL3A0Z" colab_type="code" colab={}
x_ds = tf.data.Dataset.from_tensor_slices(x_lin)
y_ds = tf.data.Dataset.from_tensor_slices(y_lin)


# + id="qgRtWDOkA-WG" colab_type="code" colab={}
def add_noise(x):
    return tf.random.normal(x.shape, mean=x)


# + id="fArSTK363LJ8" colab_type="code" colab={}
x_noisy = x_ds.map(add_noise)

# + id="J8bJ4NTb4jOh" colab_type="code" colab={}
xy_ds = tf.data.Dataset.zip((x_noisy, y_ds))

# + id="zo-VPie33w0C" colab_type="code" colab={}
model = tf.keras.Sequential([
     tf.keras.layers.Dense(10, input_shape=[1]),
     tf.keras.layers.Dense(1)]
)

# + id="NzpuWJbZ4Ugj" colab_type="code" colab={}
model.compile(loss="mse")

# + id="B11CCqfjAdC9" colab_type="code" colab={}
xy_ds.take(1)

# + id="p3zUJMrf4btG" colab_type="code" colab={}
model.fit(xy_ds.batch(5).repeat(), epochs=100, steps_per_epoch=10)

# + id="90R79fP-NEpj" colab_type="code" colab={}
model.predict(xy_ds.take(1), steps=1)

# + id="5nmlVfZXxDjc" colab_type="code" colab={}


# + [markdown] id="82VJJrMKClxf" colab_type="text"
# # Numpy (python) wrap

# + [markdown] id="aP7_kHAcFyec" colab_type="text"
# Sometimes it will be easier to repurpose previously written Python code, most likely a function using numpy or scipy, than rewriting it with tensorflow methods. The dataset objects are fundamentally built for graph mode, even if there are ways to expose pieces with eager execution, and as such datasets want to operate and use tensors. 

# + id="yWQHpCydCoSq" colab_type="code" colab={}
def add_random_noise(x, mu=0.0, std=1.0):
    if isinstance(x, np.float):
        x_noisy = x + std * np.random.randn(1) + mu
    else:
        x_noisy = x + std * np.random.randn(*x.shape) + mu
    return np.float32(x_noisy)


# + id="a82VTaOpIPfJ" colab_type="code" colab={}
arr = np.array([7, 2, 1, 6, 3, 5, 9], dtype=np.float32)
dataset = tf.data.Dataset.from_tensor_slices(arr)


# + id="Bmj3PQFLI9tc" colab_type="code" colab={}
@tf.function(input_signature=[tf.TensorSpec(None, tf.float32)]) 
def tf_add_random(input):
    mu = 0.
    sigma = 1.
    np_random = lambda x: add_random_noise(x, mu, sigma)
    y = tf.numpy_function(np_random, [input], tf.float32) 
    return y


# + id="0dD5GTBSJNN3" colab_type="code" colab={}
list(dataset.map(tf_add_random).as_numpy_iterator())

StarcoderdataJupyterScriptsDedupFiltered