# Path: rigl/cifar_resnet/resnet_train_eval.py
# FLAGS = flags.FLAGS
# PARAM_SUFFIXES = ('gamma', 'beta', 'weights', 'biases')
# MASK_SUFFIX = 'mask'
# CLASSES = [
#     'airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse',
#     'ship', 'truck'
# ]
# def create_eval_metrics(labels, logits):
# def train_fn(training_method, global_step, total_loss, train_dir, accuracy,
#              top_5_accuracy):
# def build_model(mode,
#                 images,
#                 labels,
#                 training_method='baseline',
#                 num_classes=10,
#                 depth=10,
#                 width=4):
# def wide_resnet_w_pruning(features, labels, mode, params):
#     def init_fn(scaffold, session):
# def main(argv):
# 
# Path: rigl/cifar_resnet/data_helper.py
# def input_fn(params):
#   """Provides batches of CIFAR data.
# 
#   Args:
#     params: A dictionary with a set of arguments, namely:
#       * batch_size (int32), specifies data points in a batch
#       * data_split (string), designates train or eval
#       * data_dictionary (string), specifies directory location of input dataset
# 
#   Returns:
#     images: A float32`Tensor` of size [batch_size, 32, 32, 3].
#     labels: A  int32`Tensor` of size [batch_size, num_classes].
#   """
# 
#   def parse_serialized_example(record):
#     """Parses a CIFAR10 example."""
#     image = record['image']
#     label = tf.cast(record['label'], tf.int32)
#     image = tf.cast(image, tf.float32)
#     image = tf.image.per_image_standardization(image)
#     if data_split == 'train':
#       image = preprocess_train(image, IMG_SIZE, IMG_SIZE)
#     return image, label
# 
#   data_split = params['data_split']
#   batch_size = params['batch_size']
#   if data_split == 'eval':
#     data_split = 'test'
#   dataset = tfds.load('cifar10:3.*.*', split=data_split)
# 
#   # we only repeat an example and shuffle inputs during training
#   if data_split == 'train':
#     dataset = dataset.repeat().shuffle(buffer_size=50000)
# 
#   # deserialize record into tensors and apply pre-processing.
#   dataset = dataset.map(parse_serialized_example).prefetch(batch_size)
# 
#   # at test time, for the final batch we drop remaining examples so that no
#   # example is seen twice.
#   dataset = dataset.batch(batch_size)
# 
#   images_batch, labels_batch = tf.data.make_one_shot_iterator(
#       dataset).get_next()
# 
#   return (tf.reshape(images_batch, [batch_size, IMG_SIZE, IMG_SIZE, 3]),
#           tf.reshape(labels_batch, [batch_size]))


import os
import absl.testing.parameterized as parameterized
import tensorflow.compat.v1 as tf
from absl import flags
from absl import logging
from rigl.cifar_resnet import resnet_train_eval
from rigl.cifar_resnet.data_helper import input_fn
from tensorflow.contrib.model_pruning.python import pruning

    tf.reset_default_graph()
    g = tf.Graph()
    with g.as_default():
      test_inputs, test_labels = self.get_next()

      with self.test_session() as sess:
        test_images_out, test_labels_out = sess.run([test_inputs, test_labels])
        self.assertAllEqual(test_images_out.shape, [BATCH_SIZE, 32, 32, 3])
        self.assertAllEqual(test_labels_out.shape, [BATCH_SIZE])

  @parameterized.parameters(
      {
          'training_method': 'baseline',
      },
      {
          'training_method': 'threshold',
      },
      {
          'training_method': 'rigl',
      },
  )
  def testTrainingStep(self, training_method):

    tf.reset_default_graph()
    g = tf.Graph()
    with g.as_default():

      images, labels = self.get_next()


# Path: rigl/cifar_resnet/resnet_train_eval.py
# FLAGS = flags.FLAGS
# PARAM_SUFFIXES = ('gamma', 'beta', 'weights', 'biases')
# MASK_SUFFIX = 'mask'
# CLASSES = [
#     'airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse',
#     'ship', 'truck'
# ]
# def create_eval_metrics(labels, logits):
# def train_fn(training_method, global_step, total_loss, train_dir, accuracy,
#              top_5_accuracy):
# def build_model(mode,
#                 images,
#                 labels,
#                 training_method='baseline',
#                 num_classes=10,
#                 depth=10,
#                 width=4):
# def wide_resnet_w_pruning(features, labels, mode, params):
#     def init_fn(scaffold, session):
# def main(argv):
# 
# Path: rigl/cifar_resnet/data_helper.py
# def input_fn(params):
#   """Provides batches of CIFAR data.
# 
#   Args:
#     params: A dictionary with a set of arguments, namely:
#       * batch_size (int32), specifies data points in a batch
#       * data_split (string), designates train or eval
#       * data_dictionary (string), specifies directory location of input dataset
# 
#   Returns:
#     images: A float32`Tensor` of size [batch_size, 32, 32, 3].
#     labels: A  int32`Tensor` of size [batch_size, num_classes].
#   """
# 
#   def parse_serialized_example(record):
#     """Parses a CIFAR10 example."""
#     image = record['image']
#     label = tf.cast(record['label'], tf.int32)
#     image = tf.cast(image, tf.float32)
#     image = tf.image.per_image_standardization(image)
#     if data_split == 'train':
#       image = preprocess_train(image, IMG_SIZE, IMG_SIZE)
#     return image, label
# 
#   data_split = params['data_split']
#   batch_size = params['batch_size']
#   if data_split == 'eval':
#     data_split = 'test'
#   dataset = tfds.load('cifar10:3.*.*', split=data_split)
# 
#   # we only repeat an example and shuffle inputs during training
#   if data_split == 'train':
#     dataset = dataset.repeat().shuffle(buffer_size=50000)
# 
#   # deserialize record into tensors and apply pre-processing.
#   dataset = dataset.map(parse_serialized_example).prefetch(batch_size)
# 
#   # at test time, for the final batch we drop remaining examples so that no
#   # example is seen twice.
#   dataset = dataset.batch(batch_size)
# 
#   images_batch, labels_batch = tf.data.make_one_shot_iterator(
#       dataset).get_next()
# 
#   return (tf.reshape(images_batch, [batch_size, IMG_SIZE, IMG_SIZE, 3]),
#           tf.reshape(labels_batch, [batch_size]))


# Path: rigl/cifar_resnet/data_helper_test.py
import os
import absl.testing.parameterized as parameterized
import tensorflow.compat.v1 as tf
from absl import flags
from absl import logging
from rigl.cifar_resnet import resnet_train_eval
from rigl.cifar_resnet.data_helper import input_fn
from tensorflow.contrib.model_pruning.python import pruning


    tf.reset_default_graph()
    g = tf.Graph()
    with g.as_default():
      test_inputs, test_labels = self.get_next()

      with self.test_session() as sess:
        test_images_out, test_labels_out = sess.run([test_inputs, test_labels])
        self.assertAllEqual(test_images_out.shape, [BATCH_SIZE, 32, 32, 3])
        self.assertAllEqual(test_labels_out.shape, [BATCH_SIZE])

  @parameterized.parameters(
      {
          'training_method': 'baseline',
      },
      {
          'training_method': 'threshold',
      },
      {
          'training_method': 'rigl',
      },
  )
  def testTrainingStep(self, training_method):

    tf.reset_default_graph()
    g = tf.Graph()
    with g.as_default():

      images, labels = self.get_next()


      global_step, _, _, logits = resnet_train_eval.build_model(