research/improve_nas/trainer/cifar10.py [96:133]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    x = None
    y = None
    if partition == 'train':
      x, y = x_train, y_train
    else:
      x, y = x_test, y_test

    dataset = tf.data.Dataset.from_tensor_slices((x, y.astype(np.int32)))
    return dataset.cache()

  def _shape(self):
    """Returns a 3-dimensional list with the shape of the image."""
    return [32, 32, 3]

  def get_input_fn(self,
                   partition,
                   mode,
                   batch_size,
                   preprocess=True,
                   use_tpu=False):
    """See `data.Provider` get_input_fn."""

    def input_fn(params=None):
      """Provides batches of CIFAR images.

      Args:
        params: A dict containing the batch_size on TPU, otherwise None.

      Returns:
        images: A `Tensor` of size [batch_size, 32, 32, 3]
        labels: A `Tensor` of size [batch_size, 1],
      """

      batch_size_ = batch_size
      if use_tpu:
        batch_size_ = params.get('batch_size', batch_size)

      training = mode == tf.estimator.ModeKeys.TRAIN
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -



research/improve_nas/trainer/cifar100.py [100:137]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    x = None
    y = None
    if partition == 'train':
      x, y = x_train, y_train
    else:
      x, y = x_test, y_test

    dataset = tf.data.Dataset.from_tensor_slices((x, y.astype(np.int32)))
    return dataset.cache()

  def _shape(self):
    """Returns a 3-dimensional list with the shape of the image."""
    return [32, 32, 3]

  def get_input_fn(self,
                   partition,
                   mode,
                   batch_size,
                   preprocess=True,
                   use_tpu=False):
    """See `data.Provider` get_input_fn."""

    def input_fn(params=None):
      """Provides batches of CIFAR images.

      Args:
        params: A dict containing the batch_size on TPU, otherwise None.

      Returns:
        images: A `Tensor` of size [batch_size, 32, 32, 3]
        labels: A `Tensor` of size [batch_size, 1],
      """

      batch_size_ = batch_size
      if use_tpu:
        batch_size_ = params.get('batch_size', batch_size)

      training = mode == tf.estimator.ModeKeys.TRAIN
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -