9de6656a779e73ac61995bd87044af21b3f37951,models/experimental/amoeba_net/amoeba_net_model.py,InputPipeline,input_fn,#InputPipeline#Any#,454

Before Change

      images, labels = dataset.make_one_shot_iterator().get_next()
    else:
      images = tf.random_uniform(
          [batch_size, self.hparams.image_size,
           self.hparams.image_size, 3], minval=-1, maxval=1)
      labels = tf.random_uniform(
          [batch_size], minval=0, maxval=999, dtype=tf.int32)
    if self.hparams.transpose_enabled:
      images = tensor_transform_fn(images, params["output_perm"])
    one_hot_labels = tf.one_hot(labels, self.num_classes)
    return images, one_hot_labels


def tensor_transform_fn(data, perm):
  Transpose function.

After Change

    // than a full batch size. As long as this validation is done with
    // consistent batch size, exactly the same images will be used.
    if not self.is_training:
      dataset = dataset.apply(batching.filter_irregular_batches(batch_size))

    dataset = dataset.map(
        lambda images, labels: (tf.transpose(images, [1, 2, 3, 0]), labels),
        num_parallel_calls=8)

    // For XLA, we must used fixed shapes. Because we repeat the source training
    // dataset indefinitely, this is not a dangerous operation.
    //
    // When evaluating, prevent accidentally evaluating the same image twice by
    // dropping the final batch if it is less than a full batch size. As long as
    // this validation is done with consistent batch size, exactly the same
    // images will be used.
    def set_shapes(images, labels):
      images.set_shape(images.get_shape().merge_with(
          tf.TensorShape([None, None, None, batch_size])))
      labels.set_shape(labels.get_shape().merge_with(
          tf.TensorShape([batch_size])))
      return images, labels

    if self.is_training:
      dataset = dataset.map(set_shapes)

    dataset = dataset.prefetch(32)  // Prefetch overlaps in-feed with training
    return dataset  // Must return the dataset and not tensors for high perf!


class LoadEMAHook(tf.train.SessionRunHook):

Instances