TensorFlow定义Python包装器用于数据集和迭代器

由 Carrie 创建, 最后一次修改 2017-09-07

#版权所有2017 TensorFlow作者.版权所有.

#根据Apache许可证版本2.0(“许可证”)许可;

#除非符合许可证,否则您不得使用此文件.

#您可以获得许可证的副本

#http      ://www.apache.org/licenses/LICENSE-2.0

#除非适用法律要求或书面同意软件

根据许可证分发的#分发在“按原样”基础上,

#无明示或暗示的任何种类的保证或条件.

#查看有关权限的特定语言的许可证

#许可证下的限制.

# =============================================== =============================

""数据集和迭代器的Python包装器.""

from __future__ import absolute_import

from __future__ import division

from __future__ import print_function

import abc

import numpy as np

from tensorflow.contrib.data.python.framework import function

from tensorflow.python.framework import constant_op

from tensorflow.python.framework import dtypes

from tensorflow.python.framework import ops

from tensorflow.python.framework import random_seed

from tensorflow.python.framework import sparse_tensor as sparse_tensor_lib

from tensorflow.python.framework import tensor_shape

from tensorflow.python.framework import tensor_util

from tensorflow.python.ops import array_ops

from tensorflow.python.ops import control_flow_ops

from tensorflow.python.ops import gen_dataset_ops

from tensorflow.python.ops import logging_ops

from tensorflow.python.ops import math_ops

from tensorflow.python.ops import parsing_ops

from tensorflow.python.ops import random_ops

from tensorflow.python.ops import resource_variable_ops

from tensorflow.python.platform import gfile

from tensorflow.python.util import nest

class Iterator(object):

  """Represents the state of iterating through a `Dataset`."""

  def __init__(self, iterator_resource, initializer, output_types,

               output_shapes):

    """Creates a new iterator from the given iterator resource.

    NOTE(mrry): Most users will not call this initializer directly, and will

    instead use `Iterator.from_dataset()` or `Dataset.make_one_shot_iterator()`.

    Args:

      iterator_resource: A `tf.resource` scalar `tf.Tensor` representing the

        iterator.

      initializer: A `tf.Operation` that should be run to initialize this

        iterator.

      output_types: A nested structure of `tf.DType` objects corresponding to

        each component of an element of this iterator.

      output_shapes: A nested structure of `tf.TensorShape` objects

        corresponding to each component of an element of this dataset.

    """

    self._iterator_resource = iterator_resource

    self._initializer = initializer

    self._output_types = output_types

    self._output_shapes = output_shapes

  @staticmethod

  def from_dataset(dataset, shared_name=None):

    """Creates a new, uninitialized `Iterator` from the given `Dataset`.

    To initialize this iterator, you must run its `initializer`:

    ```python

    dataset = ...

    iterator = Iterator.from_dataset(dataset)

    # ...

    sess.run(iterator.initializer)

    ```

    Args:

      dataset: A `Dataset` object.

      shared_name: (Optional.) If non-empty, this iterator will be shared under

        the given name across multiple sessions that share the same devices

        (e.g. when using a remote server).

    Returns:

      An `Iterator`.

    """

    if shared_name is None:

      shared_name = ""

    iterator_resource = gen_dataset_ops.iterator(

        container="",

        shared_name=shared_name,

        output_types=nest.flatten(dataset.output_types),

        output_shapes=nest.flatten(dataset.output_shapes))

    initializer = gen_dataset_ops.make_iterator(dataset.make_dataset_resource(),

                                                iterator_resource)

    return Iterator(iterator_resource, initializer, dataset.output_types,

                    dataset.output_shapes)

  @staticmethod

  def from_structure(output_types, output_shapes=None, shared_name=None):

    """Creates a new, uninitialized `Iterator` with the given structure.

    This iterator-constructing method can be used to create an iterator that

    is reusable with many different datasets.

    The returned iterator is not bound to a particular dataset, and it has

    no `initializer`. To initialize the iterator, run the operation returned by

    `Iterator.make_initializer(dataset)`.

    The following is an example

    ```python

    iterator = Iterator.from_structure(tf.int64, tf.TensorShape([]))

    dataset_range = Dataset.range(10)

    range_initializer = iterator.make_initializer(dataset_range)

    dataset_evens = dataset_range.filter(lambda x: x % 2 == 0)

    evens_initializer = iterator.make_initializer(dataset_evens)

    # Define a model based on the iterator; in this example, the model_fn

    # is expected to take scalar tf.int64 Tensors as input (see

    # the definition of 'iterator' above).

    prediction, loss = model_fn(iterator.get_next())

    # Train for `num_epochs`, where for each epoch, we first iterate over

    # dataset_range, and then iterate over dataset_evens.

    for _ in range(num_epochs):

      # Initialize the iterator to `dataset_range`

      sess.run(range_initializer)

      while True:

        try:

          pred, loss_val = sess.run([prediction, loss])

        except tf.errors.OutOfRangeError:

          break

      # Initialize the iterator to `dataset_evens`

      sess.run(evens_initializer)

      while True:

        try:

          pred, loss_val = sess.run([prediction, loss])

        except tf.errors.OutOfRangeError:

          break

    ```

    Args:

      output_types: A nested structure of `tf.DType` objects corresponding to

        each component of an element of this iterator.

      output_shapes: (Optional.) A nested structure of `tf.TensorShape` objects

        corresponding to each component of an element of this dataset. If

        omitted, each component will have an unconstrainted shape.

      shared_name: (Optional.) If non-empty, this iterator will be shared under

        the given name across multiple sessions that share the same devices

        (e.g. when using a remote server).

    Returns:

      An `Iterator`.

    Raises:

      TypeError: If the structures of `output_shapes` and `output_types` are

        not the same.

    """

    output_types = nest.map_structure(dtypes.as_dtype, output_types)

    if output_shapes is None:

      output_shapes = nest.map_structure(

          lambda _: tensor_shape.TensorShape(None), output_types)

    else:

      output_shapes = nest.map_structure_up_to(

          output_types, tensor_shape.as_shape, output_shapes)

    nest.assert_same_structure(output_types, output_shapes)

    if shared_name is None:

      shared_name = ""

    iterator_resource = gen_dataset_ops.iterator(

        container="",

        shared_name=shared_name,

        output_types=nest.flatten(output_types),

        output_shapes=nest.flatten(output_shapes))

    return Iterator(iterator_resource, None, output_types, output_shapes)

  @property

  def initializer(self):

    """A `tf.Operation` that should be run to initialize this iterator.

    Returns:

      A `tf.Operation` that should be run to initialize this iterator

    Raises:

      ValueError: If this iterator initializes itself automatically.

    """

    if self._initializer is not None:

      return self._initializer

    else:

      # TODO(mrry): Consider whether one-shot iterators should have

      # initializers that simply reset their state to the beginning.

      raise ValueError("Iterator does not have an initializer.")

  def make_initializer(self, dataset):

    """Returns a `tf.Operation` that initializes this iterator on `dataset`.

    Args:

      dataset: A `Dataset` with compatible structure to this iterator.

    Returns:

      A `tf.Operation` that can be run to initialize this iterator on the given

      `dataset`.

    Raises:

      TypeError: If `dataset` and this iterator do not have a compatible

        element structure.

    """

    nest.assert_same_structure(self._output_types, dataset.output_types)

    nest.assert_same_structure(self._output_shapes, dataset.output_shapes)

    for iterator_dtype, dataset_dtype in zip(

        nest.flatten(self._output_types), nest.flatten(dataset.output_types)):

      if iterator_dtype != dataset_dtype:

        raise TypeError(

            "Expected output types %r but got dataset with output types %r." %

            (self._output_types, dataset.output_types))

    for iterator_shape, dataset_shape in zip(

        nest.flatten(self._output_shapes), nest.flatten(dataset.output_shapes)):

      if not iterator_shape.is_compatible_with(dataset_shape):

        raise TypeError("Expected output shapes compatible with %r but got "

                        "dataset with output shapes %r." %

                        (self._output_shapes, dataset.output_shapes))

    return gen_dataset_ops.make_iterator(dataset.make_dataset_resource(),

                                         self._iterator_resource)

  def get_next(self, name=None):

    """Returns a nested structure of `tf.Tensor`s containing the next element.

    Args:

      name: (Optional.) A name for the created operation.

    Returns:

      A nested structure of `tf.Tensor` objects.

    """

    return nest.pack_sequence_as(

        self._output_types,

        gen_dataset_ops.iterator_get_next(

            self._iterator_resource,

            output_types=nest.flatten(self._output_types),

            output_shapes=nest.flatten(self._output_shapes),

            name=name))

  def dispose_op(self, name=None):

    """Returns a `tf.Operation` that destroys this iterator.

    The returned operation may be used to release any resources consumed by

    this iterator without closing the session.

    Args:

      name: (Optional.) A name for the created operation.

    Returns:

      A `tf.Operation`.

    """

    return gen_dataset_ops.iterator_dispose(self._iterator_resource, name=name)

  @property

  def output_shapes(self):

    """Returns the shape of each component of an element of this iterator.

    Returns:

      A nested structure of `tf.TensorShape` objects corresponding to each

      component of an element of this iterator.

    """

    return self._output_shapes

  @property

  def output_types(self):

    """Returns the type of each component of an element of this iterator.

    Returns:

      A nested structure of `tf.DType` objects corresponding to each component

      of an element of this iterator.

    """

    return self._output_types

def _calculate_acceptance_probs(initial_probs, target_probs):

  """Calculate the per-class acceptance rates.

  Args:

    initial_probs: The class probabilities of the data.

    target_probs: The desired class proportion in minibatches.

  Returns:

    A list of the per-class acceptance probabilities.

  This method is based on solving the following analysis:

  Let F be the probability of a rejection (on any example).

  Let p_i be the proportion of examples in the data in class i (init_probs)

  Let a_i is the rate the rejection sampler should *accept* class i

  Let t_i is the target proportion in the minibatches for class i (target_probs)

  ```

  F = sum_i(p_i * (1-a_i))

    = 1 - sum_i(p_i * a_i)     using sum_i(p_i) = 1

  ```

  An example with class `i` will be accepted if `k` rejections occur, then an

  example with class `i` is seen by the rejector, and it is accepted. This can

  be written as follows:

  ```

  t_i = sum_k=0^inf(F^k * p_i * a_i)

      = p_i * a_j / (1 - F)    using geometric series identity, since 0 <= F < 1

      = p_i * a_i / sum_j(p_j * a_j)        using F from above

  ```

  Note that the following constraints hold:

  ```

  0 <= p_i <= 1, sum_i(p_i) = 1

  0 <= a_i <= 1

  0 <= t_i <= 1, sum_i(t_i) = 1

  ```

  A solution for a_i in terms of the other variabes is the following:

    ```a_i = (t_i / p_i) / max_i[t_i / p_i]```

  """

  # Add tiny to initial_probs to avoid divide by zero.

  denom = (initial_probs + np.finfo(initial_probs.dtype.as_numpy_dtype).tiny)

  ratio_l = target_probs / denom

  # Calculate list of acceptance probabilities.

  max_ratio = math_ops.reduce_max(ratio_l)

  return ratio_l / max_ratio

def _estimate_data_distribution(c, num_examples_per_class_seen):

  """Estimate data distribution as labels are seen.

  Args:

    c: The class labels.  Type `int32`, shape `[batch_size]`.

    num_examples_per_class_seen: A `ResourceVariable` containing counts.

      Type `int64`, shape `[num_classes]`.

  Returns:

    dist: The updated distribution.  Type `float32`, shape `[num_classes]`.

  """

  num_classes = num_examples_per_class_seen.get_shape()[0].value

  # Update the class-count based on what labels are seen in

  # batch.  But do this asynchronously to avoid performing a

  # cross-device round-trip.  Just use the cached value.

  num_examples_per_class_seen = num_examples_per_class_seen.assign_add(

      math_ops.reduce_sum(

          array_ops.one_hot(c, num_classes, dtype=dtypes.int64),

          0))

  init_prob_estimate = math_ops.truediv(

      num_examples_per_class_seen,

      math_ops.reduce_sum(num_examples_per_class_seen))

  return math_ops.cast(init_prob_estimate, dtypes.float32)

class Dataset(object):

  """Represents a potentially large set of elements.

  A `Dataset` can be used to represent an input pipeline as a

  collection of elements (nested structures of tensors) and a "logical

  plan" of transformations that act on those elements.

  """

  __metaclass__ = abc.ABCMeta

  def __init__(self):

    pass

  @abc.abstractmethod

  def make_dataset_resource(self):

    """Creates a `tf.Tensor` of  `tf.resource` tensor representing this dataset.

    Returns:

      A scalar `tf.Tensor` of `tf.resource` type, which represents this dataset.

    """

    raise NotImplementedError("Dataset.make_dataset_resource")

  def make_initializable_iterator(self, shared_name=None):

    """Creates an `Iterator` for enumerating the elements of this dataset.

    **N.B.** The returned iterator will be in an uninitialized state,

    and you must run the `iterator.initializer` operation before using it.

    Args:

      shared_name: (Optional.) If non-empty, this iterator will be shared under

        the given name across multiple sessions that share the same devices

        (e.g. when using a remote server).

    Returns:

      An `Iterator` over the elements of this dataset.

    """

    return Iterator.from_dataset(self, shared_name)

  def make_one_shot_iterator(self):

    """Creates an `Iterator` for enumerating the elements of this dataset.

    **N.B.** The returned iterator will be initialized automatically.

    A "one-shot" iterator does not currently support re-initialization.

    Returns:

      An `Iterator` over the elements of this dataset.

    """

    # NOTE(mrry): We capture by value here to ensure that `_make_dataset()` is

    # a 0-argument function.

    @function.Defun(capture_by_value=True)

    def _make_dataset():

      return self.make_dataset_resource()

    _make_dataset.add_to_graph(ops.get_default_graph())

    return Iterator(

        gen_dataset_ops.one_shot_iterator(

            dataset_factory=_make_dataset,

            output_types=nest.flatten(self.output_types),

            output_shapes=nest.flatten(self.output_shapes)), None,

        self.output_types, self.output_shapes)

  @abc.abstractproperty

  def output_shapes(self):

    """Returns the shape of each component of an element of this dataset.

    Returns:

      A nested structure of `tf.TensorShape` objects corresponding to each

      component of an element of this dataset.

    """

    raise NotImplementedError("Dataset.output_shapes")

  @abc.abstractproperty

  def output_types(self):

    """Returns the type of each component of an element of this dataset.

    Returns:

      A nested structure of `tf.DType` objects corresponding to each component

      of an element of this dataset.

    """

    raise NotImplementedError("Dataset.output_types")

  def __repr__(self):

    output_shapes = nest.map_structure(str, self.output_shapes)

    output_shapes = str(output_shapes).replace("'", "")

    output_types = nest.map_structure(repr, self.output_types)

    output_types = str(output_types).replace("'", "")

    return ("<%s shapes: %s, types: %s>"

            % (type(self).__name__, output_shapes, output_types))

  @staticmethod

  def from_tensors(tensors):

    """Creates a `Dataset` with a single element, comprising the given tensors.

    Args:

      tensors: A nested structure of tensors.

    Returns:

      A `Dataset`.

    """

    return TensorDataset(tensors)

  @staticmethod

  def from_tensor_slices(tensors):

    """Creates a `Dataset` whose elements are slices of the given tensors.

    Args:

      tensors: A nested structure of tensors, each having the same size in the

        0th dimension.

    Returns:

      A `Dataset`.

    """

    return TensorSliceDataset(tensors)

  @staticmethod

  def from_sparse_tensor_slices(sparse_tensor):

    """Splits each rank-N `tf.SparseTensor` in this dataset row-wise.

    Args:

      sparse_tensor: A `tf.SparseTensor`.

    Returns:

      A `Dataset` of rank-(N-1) sparse tensors.

    """

    return SparseTensorSliceDataset(sparse_tensor)

  @staticmethod

  def range(*args):

    """Creates a `Dataset` of a step-separated range of values.

    For example:

    ```python

    Dataset.range(5) == [0, 1, 2, 3, 4]

    Dataset.range(2, 5) == [2, 3, 4]

    Dataset.range(1, 5, 2) == [1, 3]

    Dataset.range(1, 5, -2) == []

    Dataset.range(5, 1) == []

    Dataset.range(5, 1, -2) == [5, 3]

    ```

    Args:

      *args: follow same semantics as python's xrange.

        len(args) == 1 -> start = 0, stop = args[0], step = 1

        len(args) == 2 -> start = args[0], stop = args[1], step = 1

        len(args) == 3 -> start = args[0], stop = args[1, stop = args[2]

    Returns:

      A `RangeDataset`.

    Raises:

      ValueError: if len(args) == 0.

    """

    return RangeDataset(*args)

  @staticmethod

  def zip(datasets):

    """Creates a `Dataset` by zipping together the given datasets.

    This method has similar semantics to the built-in `zip()` function

    in Python, with the main difference being that the `datasets`

    argument can be an arbitrary nested structure of `Dataset` objects.

    For example:

    ```python

    # NOTE: The following examples use `{ ... }` to represent the

    # contents of a dataset.

    a = { 1, 2, 3 }

    b = { 4, 5, 6 }

    c = { (7, 8), (9, 10), (11, 12) }

    d = { 13, 14 }

    # The nested structure of the `datasets` argument determines the

    # structure of elements in the resulting dataset.

    Dataset.zip((a, b)) == { (1, 4), (2, 5), (3, 6) }

    Dataset.zip((b, a)) == { (4, 1), (5, 2), (6, 3) }

    # The `datasets` argument may contain an arbitrary number of

    # datasets.

    Dataset.zip((a, b, c) == { (1, 4, (7, 8)),

                               (2, 5, (9, 10)),

                               (3, 6, (11, 12)) }

    # The number of elements in the resulting dataset is the same as

    # the size of the smallest dataset in `datasets`.

    Dataset.zip((a, d)) == { (1, 13), (2, 14) }

    ```

    Args:

      datasets: A nested structure of datasets.

    Returns:

      A `Dataset`.

    """

    return ZipDataset(datasets)

  @staticmethod

  def read_batch_features(file_pattern,

                          batch_size,

                          features,

                          reader,

                          reader_args=None,

                          randomize_input=True,

                          num_epochs=None,

                          capacity=10000):

    """Reads batches of Examples.

    Args:

      file_pattern: A string pattern or a placeholder with list of filenames.

      batch_size: A `tf.int64` scalar `tf.Tensor`, representing the number of

        consecutive elements of this dataset to combine in a single batch.

      features: A `dict` mapping feature keys to `FixedLenFeature` or

        `VarLenFeature` values. See `tf.parse_example`.

      reader: A function or class that can be called with a `filenames` tensor

        and (optional) `reader_args` and returns a `Dataset` of serialized

        Examples.

      reader_args: Additional arguments to pass to the reader class.

      randomize_input: Whether the input should be randomized.

      num_epochs: Integer specifying the number of times to read through the

        dataset. If None, cycles through the dataset forever.

      capacity: Capacity of the ShuffleDataset.

    Returns:

      A `Dataset`.

    """

    if isinstance(file_pattern, str):

      filenames = _get_file_names(file_pattern, randomize_input)

    else:

      filenames = file_pattern

    if reader_args:

      dataset = reader(filenames, *reader_args)

    else:

      dataset = reader(filenames)

    dataset = dataset.repeat(num_epochs)

    if randomize_input:

      dataset = dataset.shuffle(capacity)

    dataset = dataset.map(

        lambda x: _parse_example(nest.flatten(x), features)

    )

    dataset = dataset.batch(batch_size)

    return dataset

  def repeat(self, count=None):

    """Repeats this dataset `count` times.

    Args:

      count: (Optional.) A `tf.int64` scalar `tf.Tensor`, representing the

        number of times the elements of this dataset should be repeated. The

        default behavior (if `count` is `None` or `-1`) is for the elements to

        be repeated indefinitely.

    Returns:

      A `Dataset`.

    """

    return RepeatDataset(self, count)

  def enumerate(self, start=0):

    """Enumerate the elements of this dataset.  Similar to python's `enumerate`.

    For example:

    ```python

    # NOTE: The following examples use `{ ... }` to represent the

    # contents of a dataset.

    a = { 1, 2, 3 }

    b = { (7, 8), (9, 10), (11, 12) }

    # The nested structure of the `datasets` argument determines the

    # structure of elements in the resulting dataset.

    a.enumerate(start=5) == { (5, 1), (6, 2), (7, 3) }

    b.enumerate() == { (0, (7, 8)), (1, (9, 10)), (2, (11, 12)) }

    Args:

      start: A `tf.int64` scalar `tf.Tensor`, representing the start

        value for enumeration.

    Returns:

      A `Dataset`.

    """

    max_value = np.iinfo(dtypes.int64.as_numpy_dtype).max

    return Dataset.zip((Dataset.range(start, max_value), self))

  def shuffle(self, buffer_size, seed=None):

    """Randomly shuffles the elements of this dataset.

    Args:

      buffer_size: A `tf.int64` scalar `tf.Tensor`, representing the

        number of elements from this dataset from which the new

        dataset will sample.

      seed: (Optional.) A `tf.int64` scalar `tf.Tensor`, representing the

        random seed that will be used to create the distribution. See

        @{tf.set_random_seed} for behavior.

    Returns:

      A `Dataset`.

    """

    return ShuffleDataset(self, buffer_size, seed)

  def take(self, count):

    """Creates a `Dataset` with at most `count` elements from this dataset.

    Args:

      count: A `tf.int64` scalar `tf.Tensor`, representing the number of

        elements of this dataset that should be taken to form the new dataset.

        If `count` is -1, or if `count` is greater than the size of this

        dataset, the new dataset will contain all elements of this dataset.

    Returns:

      A `Dataset`.

    """

    return TakeDataset(self, count)

  def skip(self, count):

    """Creates a `Dataset` that skips `count` elements from this dataset.

    Args:

      count: A `tf.int64` scalar `tf.Tensor`, representing the number

        of elements of this dataset that should be skipped to form the

        new dataset.  If `count` is greater than the size of this

        dataset, the new dataset will contain no elements.  If `count`

        is -1, skips the entire dataset.

    Returns:

      A `Dataset`.

    """

    return SkipDataset(self, count)

  def batch(self, batch_size):

    """Combines consecutive elements of this dataset into batches.

    Args:

      batch_size: A `tf.int64` scalar `tf.Tensor`, representing the number of

        consecutive elements of this dataset to combine in a single batch.

    Returns:

      A `Dataset`.

    """

    return BatchDataset(self, batch_size)

  def padded_batch(self, batch_size, padded_shapes, padding_values=None):

    """Combines consecutive elements of this dataset into padded batches.

    Like `Dataset.dense_to_sparse_batch()`, this method combines

    multiple consecutive elements of this dataset, which might have

    different shapes, into a single element. The tensors in the

    resulting element have an additional outer dimension, and are

    padded to the respective shape in `padded_shapes`.

    Args:

      batch_size: A `tf.int64` scalar `tf.Tensor`, representing the number of

        consecutive elements of this dataset to combine in a single batch.

      padded_shapes: A nested structure of `tf.TensorShape` or

        `tf.int64` vector tensor-like objects representing the shape

        to which the respective component of each input element should

        be padded prior to batching. Any unknown dimensions

        (e.g. `tf.Dimension(None)` in a `tf.TensorShape` or `-1` in a

        tensor-like object) will be padded to the maximum size of that

        dimension in each batch.

      padding_values: (Optional.) A nested structure of scalar-shaped

        `tf.Tensor`, representing the padding values to use for the

        respective components.  Defaults are `0` for numeric types and

        the empty string for string types.

    Returns:

      A `Dataset`.

    """

    return PaddedBatchDataset(self, batch_size, padded_shapes, padding_values)

  def dense_to_sparse_batch(self, batch_size, row_shape):

    """Batches ragged elements of this dataset into `tf.SparseTensor`s.

    Like `Dataset.padded_batch()`, this method combines multiple

    consecutive elements of this dataset, which might have different

    shapes, into a single element. The resulting element has three

    components (`indices`, `values`, and `dense_shape`), which

    comprise a `tf.SparseTensor` that represents the same data. The

    `row_shape` represents the dense shape of each row in the

    resulting `tf.SparseTensor`, to which the effective batch size is

    prepended. For example:

    ```python

    # NOTE: The following examples use `{ ... }` to represent the

    # contents of a dataset.

    a = { ['a', 'b', 'c'], ['a', 'b'], ['a', 'b', 'c', 'd'] }

    a.dense_to_sparse_batch(batch_size=2, row_shape=[6]) == {

        ([[0, 0], [0, 1], [0, 2], [1, 0], [1, 1]],  # indices

         ['a', 'b', 'c', 'a', 'b'],                 # values

         [2, 6]),                                   # dense_shape

        ([[2, 0], [2, 1], [2, 2], [2, 3]],

         ['a', 'b', 'c', 'd'],

         [1, 6])

    }

    ```

    Args:

      batch_size: A `tf.int64` scalar `tf.Tensor`, representing the

        number of consecutive elements of this dataset to combine in a

        single batch.

      row_shape: A `tf.TensorShape` or `tf.int64` vector tensor-like

        object representing the equivalent dense shape of a row in the

        resulting `tf.SparseTensor`. Each element of this dataset must

        have the same rank as `row_shape`, and must have size less

        than or equal to `row_shape` in each dimension.

    Returns:

      A `Dataset`.

    """

    return DenseToSparseBatchDataset(self, batch_size, row_shape)

  def group_by_window(self, key_func, reduce_func, window_size):

    """Performs a windowed "group-by" operation on this dataset.

    This method maps each consecutive element in this dataset to a key

    using `key_func` and groups the elements by key. It then applies

    `reduce_func` to at most `window_size` elements matching the same

    key. All execpt the final window for each key will contain

    `window_size` elements; the final window may be smaller.

    Args:

      key_func: A function mapping a nested structure of tensors

        (having shapes and types defined by `self.output_shapes` and

        `self.output_types`) to a scalar `tf.int64` tensor.

      reduce_func: A function mapping a key and a dataset of up to `batch_size`

        consecutive elements matching that key to another dataset.

      window_size: A `tf.int64` scalar `tf.Tensor`, representing the number of

        consecutive elements matching the same key to combine in a single

        batch, which will be passed to `reduce_func`.

    Returns:

      A `Dataset`.

    """

    return GroupByWindowDataset(self, key_func, reduce_func, window_size)

  def map(self, map_func, num_threads=None, output_buffer_size=None):

    """Maps `map_func` across this datset.

    Args:

      map_func: A function mapping a nested structure of tensors (having

        shapes and types defined by `self.output_shapes` and

       `self.output_types`) to another nested structure of tensors.

      num_threads: (Optional.) A `tf.int32` scalar `tf.Tensor`, representing

        the number of threads to use for processing elements in parallel. If

        not specified, elements will be processed sequentially without

        buffering.

      output_buffer_size: (Optional.) A `tf.int64` scalar `tf.Tensor`,

        representing the maximum number of processed elements that will be

        buffered when processing in parallel.

    Returns:

      A `Dataset`.

    """

    return MapDataset(self, map_func, num_threads, output_buffer_size)

  def flat_map(self, map_func):

    """Maps `map_func` across this dataset and flattens the result.

    Args:

      map_func: A function mapping a nested structure of tensors (having shapes

        and types defined by `self.output_shapes` and `self.output_types`) to a

        `Dataset`.

    Returns:

      A `Dataset`.

    """

    return FlatMapDataset(self, map_func)

  def unbatch(self):

    """Splits elements of this dataset into sequences of consecutive elements.

    For example, if elements of this dataset are shaped `[B, a0, a1, ...]`,

    where `B` may vary from element to element, then for each element in

    this dataset, the unbatched dataset will contain `B` consecutive elements

    of shape `[a0, a1, ...]`.

    Returns:

      A `Dataset`.

    """

    return self.flat_map(map_func=Dataset.from_tensor_slices)

  def filter(self, predicate):

    """Filters this dataset according to `predicate`.

    Args:

      predicate: A function mapping a nested structure of tensors (having shapes

        and types defined by `self.output_shapes` and `self.output_types`) to a

        scalar `tf.bool` tensor.

    Returns:

      A `Dataset`.

    """

    return FilterDataset(self, predicate)

class TensorDataset(Dataset):

  """A `Dataset` with a single element, viz. a nested structure of tensors."""

  def __init__(self, tensors):

    """See `Dataset.from_tensors()` for details."""

    super(TensorDataset, self).__init__()

    with ops.name_scope("tensors"):

      self._tensors = nest.pack_sequence_as(tensors, [

          ops.convert_to_tensor(t, name="component_%d" % i)

          for i, t in enumerate(nest.flatten(tensors))

      ])

  def make_dataset_resource(self):

    return gen_dataset_ops.tensor_dataset(

        nest.flatten(self._tensors),

        output_shapes=nest.flatten(self.output_shapes))

  @property

  def output_shapes(self):

    return nest.pack_sequence_as(self._tensors,

                                 [t.shape for t in nest.flatten(self._tensors)])

  @property

  def output_types(self):

    return nest.pack_sequence_as(self._tensors,

                                 [t.dtype for t in nest.flatten(self._tensors)])

class TensorSliceDataset(Dataset):

  """A `Dataset` of slices from a nested structure of tensors."""

  def __init__(self, tensors):

    """See `Dataset.from_tensor_slices()` for details."""

    super(TensorSliceDataset, self).__init__()

    with ops.name_scope("tensors"):

      flat_tensors = [

          ops.convert_to_tensor(t, name="component_%d" % i)

          for i, t in enumerate(nest.flatten(tensors))

      ]

    self._tensors = nest.pack_sequence_as(tensors, flat_tensors)

    batch_dim = flat_tensors[0].get_shape()[0]

    for t in flat_tensors[1:]:

      batch_dim.assert_is_compatible_with(t.get_shape()[0])

  def make_dataset_resource(self):

    return gen_dataset_ops.tensor_slice_dataset(

        nest.flatten(self._tensors),

        output_shapes=nest.flatten(self.output_shapes))

  @property

  def output_shapes(self):

    return nest.pack_sequence_as(self._tensors, [

        tensor_shape.TensorShape(t.shape[1:])

        for t in nest.flatten(self._tensors)

    ])

  @property

  def output_types(self):

    return nest.pack_sequence_as(self._tensors,

                                 [t.dtype for t in nest.flatten(self._tensors)])

class SparseTensorSliceDataset(Dataset):

  """A `Dataset` that splits a rank-N `tf.SparseTensor` into its rows."""

  def __init__(self, sparse_tensor):

    """See `Dataset.from_sparse_tensor_slices()` for details."""

    super(SparseTensorSliceDataset, self).__init__()

    if not isinstance(sparse_tensor, sparse_tensor_lib.SparseTensor):

      raise TypeError("`sparse_tensor` must be a `tf.SparseTensor` object.")

    self._sparse_tensor = sparse_tensor

  def make_dataset_resource(self):

    return gen_dataset_ops.sparse_tensor_slice_dataset(

        self._sparse_tensor.indices, self._sparse_tensor.values,

        self._sparse_tensor.dense_shape)

  @property

  def output_shapes(self):

    indices_shape = self._sparse_tensor.indices.get_shape()

    shape_shape = self._sparse_tensor.dense_shape.get_shape()

    rank = (indices_shape[1] - 1).merge_with(shape_shape[0] - 1)

    num_values = tensor_shape.Dimension(None)

    return (tensor_shape.TensorShape([num_values, rank]),

            tensor_shape.TensorShape([num_values]), tensor_shape.TensorShape(

                [rank]))

  @property

  def output_types(self):

    return (dtypes.int64, self._sparse_tensor.dtype, dtypes.int64)

class ZipDataset(Dataset):

  """A `Dataset` that zips its inputs together."""

  def __init__(self, datasets):

    """See `Dataset.zip()` for details."""

    super(ZipDataset, self).__init__()

    self._datasets = datasets

  def make_dataset_resource(self):

    return gen_dataset_ops.zip_dataset(

        [ds.make_dataset_resource() for ds in nest.flatten(self._datasets)],

        output_shapes=[

            s

            for ds in nest.flatten(self._datasets)

            for s in nest.flatten(ds.output_shapes)

        ],

        output_types=[

            t

            for ds in nest.flatten(self._datasets)

            for t in nest.flatten(ds.output_types)

        ])

  @property

  def output_shapes(self):

    return nest.pack_sequence_as(self._datasets, [

        ds.output_shapes for ds in nest.flatten(self._datasets)])

  @property

  def output_types(self):

    return nest.pack_sequence_as(self._datasets, [

        ds.output_types for ds in nest.flatten(self._datasets)])

class RepeatDataset(Dataset):

  """A `Dataset` that repeats its input several times."""

  def __init__(self, input_dataset, count):

    """See `Dataset.repeat()` for details."""

    super(RepeatDataset, self).__init__()

    self._input_dataset = input_dataset

    if count is None:

      self._count = constant_op.constant(-1, dtype=dtypes.int64, name="count")

    else:

      self._count = ops.convert_to_tensor(count, dtype=dtypes.int64,

                                          name="count")

  def make_dataset_resource(self):

    return gen_dataset_ops.repeat_dataset(

        self._input_dataset.make_dataset_resource(),

        count=self._count,

        output_shapes=nest.flatten(self.output_shapes),

        output_types=nest.flatten(self.output_types))

  @property

  def output_shapes(self):

    return self._input_dataset.output_shapes

  @property

  def output_types(self):

    return self._input_dataset.output_types

class RangeDataset(Dataset):

  """A `Dataset` of a step separated range of values."""

  def __init__(self, *args):

    """See `Dataset.range()` for details."""

    super(RangeDataset, self).__init__()

    self._parse_args(*args)

  def _parse_args(self, *args):

    if len(args) == 1:

      self._start = self._build_tensor(0, "start")

      self._stop = args[0]

      self._step = self._build_tensor(1, "step")

    elif len(args) == 2:

      self._start = args[0]

      self._stop = args[1]

      self._step = self._build_tensor(1, "step")

    elif len(args) == 3:

      self._start = args[0]

      self._stop = args[1]

      self._step = args[2]

    else:

      raise ValueError("Invalid arguments to RangeDataset: %s" % str(args))

  def _build_tensor(self, int64_value, name):

    return constant_op.constant(int64_value, dtype=dtypes.int64, name=name)

  def make_dataset_resource(self):

    return gen_dataset_ops.range_dataset(

        start=self._start,

        stop=self._stop,

        step=self._step,

        output_shapes=nest.flatten(self.output_shapes),

        output_types=nest.flatten(self.output_types))

  @property

  def output_shapes(self):

    return tensor_shape.scalar()

  @property

  def output_types(self):

    return dtypes.int64

class ShuffleDataset(Dataset):

  """A `Dataset` that randomly shuffles the elements of its input."""

  def __init__(self, input_dataset, buffer_size, seed=None):

    """See `Dataset.shuffle()` for details."""

    super(ShuffleDataset, self).__init__()

    self._input_dataset = input_dataset

    self._buffer_size = ops.convert_to_tensor(

        buffer_size, dtype=dtypes.int64, name="buffer_size")

    seed, seed2 = random_seed.get_seed(seed)

    if seed is None:

      self._seed = constant_op.constant(0, dtype=dtypes.int64, name="seed")

    else:

      self._seed = ops.convert_to_tensor(seed, dtype=dtypes.int64, name="seed")

    if seed2 is None:

      self._seed2 = constant_op.constant(0, dtype=dtypes.int64, name="seed2")

    else:

      self._seed2 = ops.convert_to_tensor(seed2, dtype=dtypes.int64,

                                          name="seed2")

  def make_dataset_resource(self):

    return gen_dataset_ops.shuffle_dataset(

        self._input_dataset.make_dataset_resource(),

        buffer_size=self._buffer_size,

        seed=self._seed,

        seed2=self._seed2,

        output_shapes=nest.flatten(self.output_shapes),

        output_types=nest.flatten(self.output_types))

  @property

  def output_shapes(self):

    return self._input_dataset.output_shapes

  @property

  def output_types(self):

    return self._input_dataset.output_types

class TakeDataset(Dataset):

  """A `Dataset` containing the first `count` elements from its input."""

  def __init__(self, input_dataset, count):

    """See `Dataset.take()` for details."""

    super(TakeDataset, self).__init__()

    self._input_dataset = input_dataset

    self._count = ops.convert_to_tensor(count, dtype=dtypes.int64, name="count")

  def make_dataset_resource(self):

    return gen_dataset_ops.take_dataset(

        self._input_dataset.make_dataset_resource(),

        count=self._count,

        output_shapes=nest.flatten(self.output_shapes),

        output_types=nest.flatten(self.output_types))

  @property

  def output_shapes(self):

    return self._input_dataset.output_shapes

  @property

  def output_types(self):

    return self._input_dataset.output_types

class SkipDataset(Dataset):

  """A `Dataset` skipping the first `count` elements from its input."""

  def __init__(self, input_dataset, count):

    """See `Dataset.skip()` for details."""

    super(SkipDataset, self).__init__()

    self._input_dataset = input_dataset

    self._count = ops.convert_to_tensor(count, dtype=dtypes.int64, name="count")

  def make_dataset_resource(self):

    return gen_dataset_ops.skip_dataset(

        self._input_dataset.make_dataset_resource(),

        count=self._count,

        output_shapes=nest.flatten(self.output_shapes),

        output_types=nest.flatten(self.output_types))

  @property

  def output_shapes(self):

    return self._input_dataset.output_shapes

  @property

  def output_types(self):

    return self._input_dataset.output_types

class BatchDataset(Dataset):

  """A `Dataset` that batches contiguous elements from its input."""

  def __init__(self, input_dataset, batch_size):

    """See `Dataset.batch()` for details."""

    super(BatchDataset, self).__init__()

    self._input_dataset = input_dataset

    self._batch_size = batch_size

  def make_dataset_resource(self):

    return gen_dataset_ops.batch_dataset(

        self._input_dataset.make_dataset_resource(),

        batch_size=self._batch_size,

        output_shapes=nest.flatten(self.output_shapes),

        output_types=nest.flatten(self.output_types))

  @property

  def output_shapes(self):

    input_shapes = self._input_dataset.output_shapes

    return nest.pack_sequence_as(input_shapes, [

        tensor_shape.vector(None).concatenate(s)

        for s in nest.flatten(self._input_dataset.output_shapes)

    ])

  @property

  def output_types(self):

    return self._input_dataset.output_types

def _partial_shape_to_tensor(shape_like):

  try:

    # First attempt to convert the input to a shape, and return the

    # "canonical" tensor representation, which uses `-1` in place of

    # `None`.

    shape_like = tensor_shape.as_shape(shape_like)

    return ops.convert_to_tensor(

        [dim if dim is not None else -1 for dim in shape_like.as_list()],

        dtype=dtypes.int64)

  except (TypeError, ValueError):

    # The argument was not trivially convertible to a

    # `tf.TensorShape`, so fall back on the conversion to tensor

    # machinery.

    return ops.convert_to_tensor(shape_like, dtype=dtypes.int64)

def _padding_value_to_tensor(value, output_type):

  """Converts the padding value to a tensor.

  Args:

    value: The padding value.

    output_type: Its expected dtype.

  Returns:

    A scalar `Tensor`.

  Raises:

    ValueError: if the padding value is not a scalar.

    TypeError: if the padding value's type does not match `output_type`.

  """

  value = ops.convert_to_tensor(value, name="padding_value")

  if not value.shape.is_compatible_with(tensor_shape.scalar()):

    raise ValueError(

        "Padding value should be a scalar, but is not: %s" % value)

  if value.dtype != output_type:

    raise TypeError(

        "Padding value tensor (%s) does not match output type: %s"

        % (value, output_type))

  return value

class PaddedBatchDataset(Dataset):

  """A `Dataset` that batches and pads contiguous elements from its input."""

  def __init__(self, input_dataset, batch_size, padded_shapes, padding_values):

    """See `Dataset.batch()` for details."""

    super(PaddedBatchDataset, self).__init__()

    self._input_dataset = input_dataset

    self._batch_size = batch_size

    padding_values = (padding_values if padding_values is not None else

                      self._default_padding(input_dataset))

    self._padded_shapes = nest.map_structure_up_to(input_dataset.output_shapes,

                                                   _partial_shape_to_tensor,

                                                   padded_shapes)

    self._padding_values = nest.map_structure_up_to(input_dataset.output_shapes,

                                                    _padding_value_to_tensor,

                                                    padding_values,

                                                    input_dataset.output_types)

  def _default_padding(self, input_dataset):

    def make_zero(t):

      if t.base_dtype == dtypes.string:

        return ""

      else:

        return np.zeros_like(t.as_numpy_dtype())

    return nest.map_structure(make_zero, input_dataset.output_types)

  def make_dataset_resource(self):

    return gen_dataset_ops.padded_batch_dataset(

        self._input_dataset.make_dataset_resource(),

        batch_size=self._batch_size,

        padded_shapes=[

            ops.convert_to_tensor(s, dtype=dtypes.int64)

            for s in nest.flatten(self._padded_shapes)

        ],

        padding_values=nest.flatten(self._padding_values),

        output_shapes=nest.flatten(self.output_shapes))

  @property

  def output_shapes(self):

    def _padded_shape_to_batch_shape(s):

      return tensor_shape.vector(None).concatenate(

          tensor_util.constant_value_as_shape(s))

    return nest.map_structure(_padded_shape_to_batch_shape, self._padded_shapes)

  @property

  def output_types(self):

    return self._input_dataset.output_types

class DenseToSparseBatchDataset(Dataset):

  """A `Dataset` that batches ragged dense elements into `tf.SparseTensor`s."""

  def __init__(self, input_dataset, batch_size, row_shape):

    """See `Dataset.dense_to_sparse_batch()` for more details."""

    super(DenseToSparseBatchDataset, self).__init__()

    if not isinstance(input_dataset.output_types, dtypes.DType):

      raise TypeError("DenseToSparseDataset requires an input whose elements "

                      "have a single component, whereas the input has %r."

                      % input_dataset.output_types)

    self._input_dataset = input_dataset

    self._batch_size = batch_size

    self._row_shape = _partial_shape_to_tensor(row_shape)

  def make_dataset_resource(self):

    return gen_dataset_ops.dense_to_sparse_batch_dataset(

        self._input_dataset.make_dataset_resource(),

        self._batch_size,

        self._row_shape,

        output_shapes=self.output_shapes,

        output_types=self.output_types)

  @property

  def output_shapes(self):

    num_elements = tensor_shape.Dimension(None)

    return (tensor_shape.matrix(num_elements, self._row_shape.shape[0] + 1),

            tensor_shape.vector(num_elements),

            tensor_shape.vector(self._row_shape.shape[0] + 1))

  @property

  def output_types(self):

    return (dtypes.int64, self._input_dataset.output_types, dtypes.int64)

class _ResourceDataset(Dataset):

  """A Dataset wrapper for a tf.resource-typed function argument."""

  def __init__(self, dataset_resource, output_types, output_shapes):

    super(_ResourceDataset, self).__init__()

    self._dataset_resource = dataset_resource,

    self._output_types = output_types

    self._output_shapes = output_shapes

  def make_dataset_resource(self):

    return self._dataset_resource

  @property

  def output_shapes(self):

    return self._output_shapes

  @property

  def output_types(self):

    return self._output_types

class GroupByWindowDataset(Dataset):

  """A `Dataset` that groups its input and performs a windowed reduction."""

  def __init__(self, input_dataset, key_func, reduce_func, window_size):

    """See `Dataset.group_by_window()` for details."""

    super(GroupByWindowDataset, self).__init__()

    self._input_dataset = input_dataset

    self._window_size = window_size

    @function.Defun(*nest.flatten(input_dataset.output_types))

    def tf_key_func(*args):

      """A wrapper for Defun that facilitates shape inference."""

      # Pass in shape information from the input_dataset.

      for arg, shape in zip(args, nest.flatten(input_dataset.output_shapes)):

        arg.set_shape(shape)

      nested_args = nest.pack_sequence_as(input_dataset.output_types, args)

      if nest.is_sequence(nested_args):

        ret = key_func(*nested_args)

      else:

        ret = key_func(nested_args)

      ret = ops.convert_to_tensor(ret, dtype=dtypes.int64)

      if ret.dtype != dtypes.int64:

        raise ValueError("`key_func` must return a single tf.int64 tensor.")

      return ret

    self._key_func = tf_key_func

    self._key_func.add_to_graph(ops.get_default_graph())

    @function.Defun(dtypes.int64, dtypes.resource)

    def tf_reduce_func(key, window_dataset_resource):

      """A wrapper for Defun that facilitates shape inference."""

      key.set_shape([])

      window_dataset = _ResourceDataset(window_dataset_resource,

                                        input_dataset.output_types,

                                        input_dataset.output_shapes)

      output_dataset = reduce_func(key, window_dataset)

      if not isinstance(output_dataset, Dataset):

        raise TypeError("`reduce_func` must return a `Dataset` object.")

      self._output_types = output_dataset.output_types

      self._output_shapes = output_dataset.output_shapes

      return output_dataset.make_dataset_resource()

    self._reduce_func = tf_reduce_func

    self._reduce_func.add_to_graph(ops.get_default_graph())

  def make_dataset_resource(self):

    return gen_dataset_ops.group_by_window_dataset(

        self._input_dataset.make_dataset_resource(),

        self._key_func.captured_inputs,

        self._reduce_func.captured_inputs,

        self._window_size,

        key_func=self._key_func,

        reduce_func=self._reduce_func,

        output_types=nest.flatten(self.output_types),

        output_shapes=nest.flatten(self.output_shapes))

  @property

  def output_shapes(self):

    return self._output_shapes

  @property

  def output_types(self):

    return self._output_types

def _most_specific_compatible_shape(s1, s2):

  """Returns the most specific shape compatible with `s1` and `s2`."""

  if s1.dims is None:

    return s1

  if s2.dims is None:

    return s2

  s1.assert_same_rank(s2)

  dims = []

  for dim1, dim2 in zip(s1, s2):

    if dim1.value is None or dim2.value is None or dim1.value != dim2.value:

      dims.append(tensor_shape.Dimension(None))

    else:

      dims.append(dim1.value)

  return tensor_shape.TensorShape(dims)

class MapDataset(Dataset):

  """A `Dataset` that maps a function over elements in its input."""

  def __init__(self,

               input_dataset,

               map_func,

               num_threads=None,

               output_buffer_size=None):

    """See `Dataset.map()` for details."""

    super(MapDataset, self).__init__()

    self._input_dataset = input_dataset

    self._output_shapes = None

    self._output_types = None

    @function.Defun(*nest.flatten(input_dataset.output_types))

    def tf_map_func(*args):

      """A wrapper for Defun that facilitates shape inference."""

      # Pass in shape information from the input_dataset.

      for arg, shape in zip(args, nest.flatten(input_dataset.output_shapes)):

        arg.set_shape(shape)

      nested_args = nest.pack_sequence_as(input_dataset.output_types, args)

      if nest.is_sequence(nested_args):

        ret = map_func(*nested_args)

      else:

        ret = map_func(nested_args)

      # Extract shape information from the returned values.

      flattened_ret = [ops.convert_to_tensor(t) for t in nest.flatten(ret)]

      self._output_shapes = nest.pack_sequence_as(

          ret, [t.get_shape() for t in flattened_ret])

      self._output_types = nest.pack_sequence_as(

          ret, [t.dtype for t in flattened_ret])

      return flattened_ret

    self._map_func = tf_map_func

    self._map_func.add_to_graph(ops.get_default_graph())

    if num_threads is not None:

      self._num_threads = ops.convert_to_tensor(

          num_threads, dtype=dtypes.int32, name="num_threads")

      if output_buffer_size is not None:

        self._output_buffer_size = ops.convert_to_tensor(

            output_buffer_size, dtype=dtypes.int64, name="output_buffer_size")

      else:

        self._output_buffer_size = self._num_threads

    else:

      self._num_threads = None

      self._output_buffer_size = None

  def make_dataset_resource(self):

    input_resource = self._input_dataset.make_dataset_resource()

    if self._num_threads is None:

      return gen_dataset_ops.map_dataset(

          input_resource,

          self._map_func.captured_inputs,

          f=self._map_func,

          output_types=nest.flatten(self.output_types),

          output_shapes=nest.flatten(self.output_shapes))

    else:

      return gen_dataset_ops.parallel_map_dataset(

          input_resource,

          self._map_func.captured_inputs,

          f=self._map_func,

          num_threads=self._num_threads,

          output_buffer_size=self._output_buffer_size,

          output_types=nest.flatten(self.output_types),

          output_shapes=nest.flatten(self.output_shapes))

  @property

  def output_shapes(self):

    return self._output_shapes

  @property

  def output_types(self):

    return self._output_types

class FlatMapDataset(Dataset):

  """A `Dataset` that maps a function over its input and flattens the result."""

  def __init__(self,

               input_dataset,

               map_func):

    """See `Dataset.flat_map()` for details."""

    super(FlatMapDataset, self).__init__()

    self._input_dataset = input_dataset

    @function.Defun(*nest.flatten(input_dataset.output_types))

    def tf_map_func(*args):

      """A wrapper for Defun that facilitates shape inference."""

      # Pass in shape information from the input_dataset.

      for arg, shape in zip(args, nest.flatten(input_dataset.output_shapes)):

        arg.set_shape(shape)

      nested_args = nest.pack_sequence_as(input_dataset.output_types, args)

      if nest.is_sequence(nested_args):

        dataset = map_func(*nested_args)

      else:

        dataset = map_func(nested_args)

      if not isinstance(dataset, Dataset):

        raise TypeError("`map_func` must return a `Dataset` object.")

      self._output_types = dataset.output_types

      self._output_shapes = dataset.output_shapes

      return dataset.make_dataset_resource()

    self._map_func = tf_map_func

    self._map_func.add_to_graph(ops.get_default_graph())

  def make_dataset_resource(self):

    return gen_dataset_ops.flat_map_dataset(

        self._input_dataset.make_dataset_resource(),

        self._map_func.captured_inputs,

        f=self._map_func,

        output_types=nest.flatten(self.output_types),

        output_shapes=nest.flatten(self.output_shapes))

  @property

  def output_shapes(self):

    return self._output_shapes

  @property

  def output_types(self):

    return self._output_types

class FilterDataset(Dataset):

  """A `Dataset` that filters its input according to a predicate function."""

  def __init__(self, input_dataset, predicate):

    """See `Dataset.filter()` for details."""

    super(FilterDataset, self).__init__()

    self._input_dataset = input_dataset

    @function.Defun(*nest.flatten(input_dataset.output_types))

    def tf_predicate(*args):

      """A wrapper for Defun that facilitates shape inference."""

      # Pass in shape information from the input_dataset.

      for arg, shape in zip(args, nest.flatten(input_dataset.output_shapes)):

        arg.set_shape(shape)

      nested_args = nest.pack_sequence_as(input_dataset.output_types, args)

      if nest.is_sequence(nested_args):

        ret = predicate(*nested_args)

      else:

        ret = predicate(nested_args)

      ret = ops.convert_to_tensor(ret, dtype=dtypes.bool)

      if not (ret.dtype == dtypes.bool and

              ret.shape.is_compatible_with(tensor_shape.scalar())):

        raise ValueError("`predicate` must return a scalar boolean tensor.")

      return ret

    self._predicate = tf_predicate

    self._predicate.add_to_graph(ops.get_default_graph())

  def make_dataset_resource(self):

    return gen_dataset_ops.filter_dataset(

        self._input_dataset.make_dataset_resource(),

        other_arguments=self._predicate.captured_inputs,

        predicate=self._predicate,

        output_types=nest.flatten(self.output_types),

        output_shapes=nest.flatten(self.output_shapes))

  @property

  def output_shapes(self):

    return self._input_dataset.output_shapes

  @property

  def output_types(self):

    return self._input_dataset.output_types

class TextLineDataset(Dataset):

  """A `Dataset` comprising lines from one or more text files."""

  def __init__(self, filenames):

    """Creates a `TextLineDataset`.

    Args:

      filenames: A `tf.string` tensor containing one or more filenames.

    """

    super(TextLineDataset, self).__init__()

    self._filenames = ops.convert_to_tensor(

        filenames, dtype=dtypes.string, name="filenames")

  def make_dataset_resource(self):

    return gen_dataset_ops.text_line_dataset(self._filenames)

  @property

  def output_shapes(self):

    return tensor_shape.scalar()

  @property

  def output_types(self):

    return dtypes.string

class TFRecordDataset(Dataset):

  """A `Dataset` comprising records from one or more TFRecord files."""

  def __init__(self, filenames, compression_type=None):

    """Creates a `TFRecordDataset`.

    Args:

      filenames: A `tf.string` tensor containing one or more filenames.

      compression_type: A `tf.string` scalar evaluating to one of `""` (no

        compression), `"ZLIB"`, or `"GZIP"`.

    """

    super(TFRecordDataset, self).__init__()

    self._filenames = ops.convert_to_tensor(filenames, name="filenames")

    if compression_type is not None:

      self._compression_type = ops.convert_to_tensor(

          compression_type, dtype=dtypes.string, name="compression_type")

    else:

      self._compression_type = constant_op.constant("", name="compression_type")

  def make_dataset_resource(self):

    return gen_dataset_ops.tf_record_dataset(self._filenames,

                                             self._compression_type)

  @property

  def output_shapes(self):

    return tensor_shape.TensorShape([])

  @property

  def output_types(self):

    return dtypes.string

class FixedLengthRecordDataset(Dataset):

  """A `Dataset` of fixed-length records from one or more binary files."""

  def __init__(self,

               filenames,

               record_bytes,

               header_bytes=None,

               footer_bytes=None):

    """Creates a `FixedLengthRecordDataset`.

    Args:

      filenames: A `tf.string` tensor containing one or more filenames.

      record_bytes: A `tf.int64` scalar representing the number of bytes in

        each record.

      header_bytes: (Optional.) A `tf.int64` scalar representing the number of

        bytes to skip at the start of a file.

      footer_bytes: (Optional.) A `tf.int64` scalar representing the number of

        bytes to ignore at the end of a file.

    """

    super(FixedLengthRecordDataset, self).__init__()

    self._filenames = ops.convert_to_tensor(

        filenames, dtype=dtypes.string, name="filenames")

    self._record_bytes = ops.convert_to_tensor(

        record_bytes, dtype=dtypes.int64, name="record_bytes")

    if header_bytes is not None:

      self._header_bytes = ops.convert_to_tensor(

          header_bytes, dtype=dtypes.int64, name="header_bytes")

    else:

      self._header_bytes = constant_op.constant(

          0, dtype=dtypes.int64, name="header_bytes")

    if footer_bytes is not None:

      self._footer_bytes = ops.convert_to_tensor(

          footer_bytes, dtype=dtypes.int64, name="footer_bytes")

    else:

      self._footer_bytes = constant_op.constant(

          0, dtype=dtypes.int64, name="footer_bytes")

  def make_dataset_resource(self):

    return gen_dataset_ops.fixed_length_record_dataset(

        self._filenames, self._header_bytes, self._record_bytes,

        self._footer_bytes)

  @property

  def output_shapes(self):

    return tensor_shape.scalar()

  @property

  def output_types(self):

    return dtypes.string

def rejection_resample(dataset, class_func, target_dist,

                       initial_dist=None, seed=None):

  """Resamples this dataset to achieve a target class distribution.

  **NOTE** Resampling is performed via rejection sampling; some fraction

  of the input values will be dropped.

  Args:

    dataset: A `Dataset` object.

    class_func: A function mapping a nested structure of tensors (having

      shapes and types defined by `dataset.output_shapes` and

      `dataset.output_types`) to a scalar `tf.int32` tensor.  Values should

      be in `[0, num_classes)`.

    target_dist: A floating point type tensor, shaped `[num_classes].

    initial_dist: (Optional.)  A floating point type tensor, shaped

      `[num_classes]`.  If not provided, the true class distribution is

      estimated live in a streaming fashion.

    seed: (Optional.) Python integer seed for the resampler.

  Returns:

    A `Dataset`.

  """

  dist_estimation_batch_size = 32

  target_dist = ops.convert_to_tensor(target_dist, name="initial_dist")

  class_values_ds = dataset.map(class_func)

  if initial_dist is not None:

    initial_dist = ops.convert_to_tensor(

        initial_dist, name="initial_dist")

    acceptance_dist = _calculate_acceptance_probs(initial_dist, target_dist)

    initial_dist_ds = Dataset.from_tensors(initial_dist).repeat()

    acceptance_dist_ds = Dataset.from_tensors(acceptance_dist).repeat()

  else:

    num_classes = (target_dist.shape[0].value

                   or array_ops.shape(target_dist)[0])

    smoothing_constant = 10

    num_examples_per_class_seen = resource_variable_ops.ResourceVariable(

        initial_value=array_ops.fill(

            [num_classes], np.int64(smoothing_constant)),

        trainable=False,

        name="class_count",

        dtype=dtypes.int64)

    def update_estimate_and_tile(c):

      return array_ops.tile(

          array_ops.expand_dims(

              _estimate_data_distribution(c, num_examples_per_class_seen), 0),

          [dist_estimation_batch_size, 1])

    initial_dist_ds = (class_values_ds

                       .batch(dist_estimation_batch_size)

                       .map(update_estimate_and_tile)

                       .unbatch())

    acceptance_dist_ds = initial_dist_ds.map(

        lambda initial: _calculate_acceptance_probs(initial, target_dist))

  def maybe_warn_on_large_rejection(accept_dist, initial_dist):

    proportion_rejected = math_ops.reduce_sum(

        (1 - accept_dist) * initial_dist)

    return control_flow_ops.cond(

        math_ops.less(proportion_rejected, .5),

        lambda: accept_dist,

        lambda: logging_ops.Print(  # pylint: disable=g-long-lambda

            accept_dist, [proportion_rejected, initial_dist, accept_dist],

            message="Proportion of examples rejected by sampler is high: ",

            summarize=100,

            first_n=10))

  acceptance_dist_ds = (

      Dataset.zip((acceptance_dist_ds, initial_dist_ds))

      .map(maybe_warn_on_large_rejection))

  current_probabilities_ds = (Dataset

                              .zip((acceptance_dist_ds, class_values_ds))

                              .map(array_ops.gather))

  filtered_ds = (

      Dataset.zip((class_values_ds, current_probabilities_ds, dataset))

      .filter(lambda _1, p, _2: random_ops.random_uniform([], seed=seed) < p))

  return filtered_ds.map(lambda class_value, _, data: (class_value, data))

def read_batch_features(file_pattern,

                        batch_size,

                        features,

                        reader,

                        reader_args=None,

                        randomize_input=True,

                        num_epochs=None,

                        capacity=10000):

  """Reads batches of Examples.

  Example:

  ```

  serialized_examples = [

    features {

      feature { key: "age" value { int64_list { value: [ 0 ] } } }

      feature { key: "gender" value { bytes_list { value: [ "f" ] } } }

      feature { key: "kws" value { bytes_list { value: [ "code", "art" ] } } }

    },

    features {

      feature { key: "age" value { int64_list { value: [] } } }

      feature { key: "gender" value { bytes_list { value: [ "f" ] } } }

      feature { key: "kws" value { bytes_list { value: [ "sports" ] } } }

    }

  ]

  ```

  We can use arguments:

  ```

  features: {

    "age": FixedLenFeature([], dtype=tf.int64, default_value=-1),

    "gender": FixedLenFeature([], dtype=tf.string),

    "kws": VarLenFeature(dtype=tf.string),

  }

  ```

  And the expected output is:

  ```python

  {

    "age": [[0], [-1]],

    "gender": [["f"], ["f"]],

    "kws": SparseTensor(

      indices=[[0, 0], [0, 1], [1, 0]],

      values=["code", "art", "sports"]

      dense_shape=[2, 2]),

  }

  ```

  Args:

    file_pattern: List of files or patterns of file paths containing

      `Example` records. See `tf.gfile.Glob` for pattern rules.

    batch_size: An int representing the number of consecutive elements of this

      dataset to combine in a single batch.

    features: A `dict` mapping feature keys to `FixedLenFeature` or

      `VarLenFeature` values. See `tf.parse_example`.

    reader: A function or class that can be called with a `filenames` tensor

      and (optional) `reader_args` and returns a `Dataset` of serialized

      Examples.

    reader_args: Additional arguments to pass to the reader class.

    randomize_input: Whether the input should be randomized.

    num_epochs: Integer specifying the number of times to read through the

      dataset. If None, cycles through the dataset forever.

    capacity: Capacity of the ShuffleDataset. A large capacity ensures better

      shuffling but would increase memory usage and startup time.

  Returns:

    A dict from keys in features to Tensor or SparseTensor objects.

  """

  filenames = _get_file_names(file_pattern, randomize_input)

  if reader_args:

    dataset = reader(filenames, *reader_args)

  else:

    dataset = reader(filenames)

  dataset = dataset.repeat(num_epochs)

  if randomize_input:

    dataset = dataset.shuffle(capacity)

  dataset = dataset.batch(batch_size)

  dataset = dataset.map(lambda x: _parse_example(x, features))

  iterator = dataset.make_one_shot_iterator()

  outputs = iterator.get_next()

  index = 0

  result = {}

  for key in sorted(features.keys()):

    feature = features[key]

    if isinstance(feature, parsing_ops.FixedLenFeature):

      result[key] = outputs[index]

      index += 1

    else:

      result[key] = sparse_tensor_lib.SparseTensor(

          indices=outputs[index],

          values=outputs[index + 1],

          dense_shape=outputs[index + 2])

      index += 3

  return result

def _parse_example(serialized, features):

  parsed = parsing_ops.parse_example(serialized, features)

  result = []

  for key in sorted(features.keys()):

    val = parsed[key]

    if isinstance(val, sparse_tensor_lib.SparseTensor):

      result.extend([val.indices, val.values, val.dense_shape])

    else:

      result.append(val)

  return result

def _get_file_names(file_pattern, randomize_input):

  """Parse list of file names from pattern, optionally shuffled.

  Args:

    file_pattern: File glob pattern, or list of glob patterns.

    randomize_input: Whether to shuffle the order of file names.

  Returns:

    List of file names matching `file_pattern`.

  Raises:

    ValueError: If `file_pattern` is empty, or pattern matches no files.

  """

  if isinstance(file_pattern, list):

    if not file_pattern:

      raise ValueError("File pattern is empty.")

    file_names = []

    for entry in file_pattern:

      file_names.extend(gfile.Glob(entry))

  else:

    file_names = list(gfile.Glob(file_pattern))

  if not file_names:

    raise ValueError("No files match %s." % file_pattern)

  # Sort files so it will be deterministic for unit tests.

  if not randomize_input:

    file_names = sorted(file_names)

  return file_names

以上内容是否对您有帮助:
二维码
建议反馈
二维码