# Path: deep_architect/contrib/misc/datasets/cifar10_tf.py
# class Cifar10DataSet:
#     """Cifar10 data set.
#   Described by http://www.cs.toronto.edu/~kriz/cifar.html.
#   """
# 
#     def __init__(self, data_dir, subset='train', use_distortion=True):
#         self.data_dir = data_dir
#         self.subset = subset
#         self.use_distortion = use_distortion
# 
#     def get_filenames(self):
#         if self.subset in ['train', 'validation', 'eval']:
#             return [os.path.join(self.data_dir, self.subset + '.tfrecords')]
#         else:
#             raise ValueError('Invalid data subset "%s"' % self.subset)
# 
#     def parser(self, serialized_example):
#         """Parses a single tf.Example into image and label tensors."""
#         # Dimensions of the images in the CIFAR-10 dataset.
#         # See http://www.cs.toronto.edu/~kriz/cifar.html for a description of the
#         # input format.
#         features = tf.parse_single_example(serialized_example,
#                                            features={
#                                                'image':
#                                                tf.FixedLenFeature([],
#                                                                   tf.string),
#                                                'label':
#                                                tf.FixedLenFeature([], tf.int64),
#                                            })
#         image = tf.decode_raw(features['image'], tf.uint8)
#         image.set_shape([DEPTH * HEIGHT * WIDTH])
# 
#         # Reshape from [depth * height * width] to [depth, height, width].
#         image = tf.cast(
#             tf.transpose(tf.reshape(image, [DEPTH, HEIGHT, WIDTH]), [1, 2, 0]),
#             tf.float32)
#         label = tf.cast(features['label'], tf.int32)
# 
#         # # Custom preprocessing.
#         # image = self.preprocess(image)
# 
#         return image, label
# 
#     def make_batch(self, batch_size):
#         """Read the images and labels from 'filenames'."""
#         filenames = self.get_filenames()
#         # Repeat infinitely.
#         dataset = tf.data.TFRecordDataset(filenames)
# 
#         # Parse records.
#         dataset = dataset.map(self.parser, num_parallel_calls=2).cache()
# 
#         # Potentially shuffle records.
#         if self.subset == 'train':
#             # Ensure that the capacity is sufficiently large to provide good random
#             # shuffling.
#             dataset = dataset.apply(
#                 tf.data.experimental.shuffle_and_repeat(
#                     Cifar10DataSet.num_examples_per_epoch(self.subset)))
#         else:
#             dataset = dataset.repeat()
# 
#         # Batch it up.
#         dataset = dataset.map(self.preprocess, num_parallel_calls=2)
#         dataset = dataset.batch(batch_size, drop_remainder=True)
#         dataset = dataset.prefetch(batch_size)
# 
#         return dataset
# 
#     def preprocess(self, image, label):
#         """Preprocess a single image in [height, width, depth] layout."""
#         image = tf.image.per_image_standardization(image)
#         if self.subset == 'train' and self.use_distortion:
#             # Pad 4 pixels on each dimension of feature map, done in mini-batch
#             image = tf.image.resize_image_with_crop_or_pad(image, 40, 40)
#             image = tf.random_crop(image, [HEIGHT, WIDTH, DEPTH])
#             image = tf.image.random_flip_left_right(image)
#         return image, label
# 
#     @staticmethod
#     def num_examples_per_epoch(subset='train'):
#         if subset == 'train':
#             return 512
#             return 40000
#         elif subset == 'validation':
#             return 10000
#         elif subset == 'eval':
#             return 10000
#         else:
#             raise ValueError('Invalid data subset "%s"' % subset)


import os
import sys
import gc
import subprocess
import random
import logging
import tensorflow as tf
import numpy as np
import deep_architect.contrib.misc.evaluators.tensorflow.gcloud_utils as gcu
import deep_architect.core as co
import deep_architect.utils as ut
import deep_architect.helpers.tensorflow_eager_support as htfe
from deep_architect.contrib.misc.datasets.cifar10_tf import Cifar10DataSet
class KerasTPUEvaluator:

    def __init__(self,
                 data_dir,
                 tpu_name,
                 max_num_training_epochs=200,
                 stop_patience=20,
                 optimizer_type='adam',
                 batch_size=256,
                 lr_decay_method='constant',
                 init_lr=.001,
                 lr_decay_value=.97,
                 lr_num_epochs_per_decay=2.4,
                 lr_warmup_epochs=3.0,
                 weight_decay=0,
                 display_step=1,
                 log_output_to_terminal=True,
                 base_dir='./scratch',
                 delete_scratch_after_use=False,
                 epochs_between_evals=100,
                 use_tpu=True):
        self.tpu_name = tpu_name
        self.num_examples = Cifar10DataSet.num_examples_per_epoch()
        self.num_examples = 512
        self.batch_size = batch_size
        self.data_dir = data_dir
        self.steps_per_epoch = self.num_examples // self.batch_size
        self.steps_per_val_epoch = Cifar10DataSet.num_examples_per_epoch(
            subset='validation') // self.batch_size
        self.steps_per_val_epoch = 3

# Path: deep_architect/contrib/misc/datasets/cifar10_tf.py
# class Cifar10DataSet:
#     """Cifar10 data set.
#   Described by http://www.cs.toronto.edu/~kriz/cifar.html.
#   """
# 
#     def __init__(self, data_dir, subset='train', use_distortion=True):
#         self.data_dir = data_dir
#         self.subset = subset
#         self.use_distortion = use_distortion
# 
#     def get_filenames(self):
#         if self.subset in ['train', 'validation', 'eval']:
#             return [os.path.join(self.data_dir, self.subset + '.tfrecords')]
#         else:
#             raise ValueError('Invalid data subset "%s"' % self.subset)
# 
#     def parser(self, serialized_example):
#         """Parses a single tf.Example into image and label tensors."""
#         # Dimensions of the images in the CIFAR-10 dataset.
#         # See http://www.cs.toronto.edu/~kriz/cifar.html for a description of the
#         # input format.
#         features = tf.parse_single_example(serialized_example,
#                                            features={
#                                                'image':
#                                                tf.FixedLenFeature([],
#                                                                   tf.string),
#                                                'label':
#                                                tf.FixedLenFeature([], tf.int64),
#                                            })
#         image = tf.decode_raw(features['image'], tf.uint8)
#         image.set_shape([DEPTH * HEIGHT * WIDTH])
# 
#         # Reshape from [depth * height * width] to [depth, height, width].
#         image = tf.cast(
#             tf.transpose(tf.reshape(image, [DEPTH, HEIGHT, WIDTH]), [1, 2, 0]),
#             tf.float32)
#         label = tf.cast(features['label'], tf.int32)
# 
#         # # Custom preprocessing.
#         # image = self.preprocess(image)
# 
#         return image, label
# 
#     def make_batch(self, batch_size):
#         """Read the images and labels from 'filenames'."""
#         filenames = self.get_filenames()
#         # Repeat infinitely.
#         dataset = tf.data.TFRecordDataset(filenames)
# 
#         # Parse records.
#         dataset = dataset.map(self.parser, num_parallel_calls=2).cache()
# 
#         # Potentially shuffle records.
#         if self.subset == 'train':
#             # Ensure that the capacity is sufficiently large to provide good random
#             # shuffling.
#             dataset = dataset.apply(
#                 tf.data.experimental.shuffle_and_repeat(
#                     Cifar10DataSet.num_examples_per_epoch(self.subset)))
#         else:
#             dataset = dataset.repeat()
# 
#         # Batch it up.
#         dataset = dataset.map(self.preprocess, num_parallel_calls=2)
#         dataset = dataset.batch(batch_size, drop_remainder=True)
#         dataset = dataset.prefetch(batch_size)
# 
#         return dataset
# 
#     def preprocess(self, image, label):
#         """Preprocess a single image in [height, width, depth] layout."""
#         image = tf.image.per_image_standardization(image)
#         if self.subset == 'train' and self.use_distortion:
#             # Pad 4 pixels on each dimension of feature map, done in mini-batch
#             image = tf.image.resize_image_with_crop_or_pad(image, 40, 40)
#             image = tf.random_crop(image, [HEIGHT, WIDTH, DEPTH])
#             image = tf.image.random_flip_left_right(image)
#         return image, label
# 
#     @staticmethod
#     def num_examples_per_epoch(subset='train'):
#         if subset == 'train':
#             return 512
#             return 40000
#         elif subset == 'validation':
#             return 10000
#         elif subset == 'eval':
#             return 10000
#         else:
#             raise ValueError('Invalid data subset "%s"' % subset)


# Path: dev/evaluators/tpu_keras_evaluator.py
import os
import sys
import gc
import subprocess
import random
import logging
import tensorflow as tf
import numpy as np
import deep_architect.contrib.misc.evaluators.tensorflow.gcloud_utils as gcu
import deep_architect.core as co
import deep_architect.utils as ut
import deep_architect.helpers.tensorflow_eager_support as htfe
from deep_architect.contrib.misc.datasets.cifar10_tf import Cifar10DataSet

class KerasTPUEvaluator:

    def __init__(self,
                 data_dir,
                 tpu_name,
                 max_num_training_epochs=200,
                 stop_patience=20,
                 optimizer_type='adam',
                 batch_size=256,
                 lr_decay_method='constant',
                 init_lr=.001,
                 lr_decay_value=.97,
                 lr_num_epochs_per_decay=2.4,
                 lr_warmup_epochs=3.0,
                 weight_decay=0,
                 display_step=1,
                 log_output_to_terminal=True,
                 base_dir='./scratch',
                 delete_scratch_after_use=False,
                 epochs_between_evals=100,
                 use_tpu=True):
        self.tpu_name = tpu_name
        self.num_examples = Cifar10DataSet.num_examples_per_epoch()
        self.num_examples = 512
        self.batch_size = batch_size
        self.data_dir = data_dir
        self.steps_per_epoch = self.num_examples // self.batch_size
        self.steps_per_val_epoch = Cifar10DataSet.num_examples_per_epoch(
            subset='validation') // self.batch_size
        self.steps_per_val_epoch = 3

        self.steps_per_test_epoch = Cifar10DataSet.num_examples_per_epoch(