# Path: rigl/imagenet_resnet/resnet_model.py
# FLAGS = flags.FLAGS
# BATCH_NORM_DECAY = 0.9
# BATCH_NORM_EPSILON = 1e-5
# def batch_norm_relu(inputs, is_training, relu=True, init_zero=False,
#                     data_format='channels_first'):
# def fixed_padding(inputs, kernel_size, data_format='channels_first'):
#   def __init__(self, sparsity, seed=None, dtype=tf.float32):
#   def __call__(self, *args, **kwargs):
#   def get_config(self):
#   def __init__(self, sparsity, seed=None, dtype=tf.float32):
#   def __call__(self, shape, dtype=None, partition_info=None):
#   def get_config(self):
#   def __init__(self, sparsity, seed=None, dtype=tf.float32):
#   def __call__(self, shape, dtype=None, partition_info=None):
#   def get_config(self):
# def _pick_initializer(kernel_initializer, init_method, pruning_method,
#                       end_sparsity):
# def conv2d_fixed_padding(inputs,
#                          filters,
#                          kernel_size,
#                          strides,
#                          pruning_method='baseline',
#                          init_method='baseline',
#                          data_format='channels_first',
#                          end_sparsity=0.,
#                          weight_decay=0.,
#                          init_scale=1.0,
#                          name=None):
# def residual_block_(inputs,
#                     filters,
#                     is_training,
#                     strides,
#                     use_projection=False,
#                     pruning_method='baseline',
#                     init_method='baseline',
#                     data_format='channels_first',
#                     end_sparsity=0.,
#                     weight_decay=0.,
#                     name=''):
# def bottleneck_block_(inputs,
#                       filters,
#                       is_training,
#                       strides,
#                       use_projection=False,
#                       pruning_method='baseline',
#                       init_method='baseline',
#                       data_format='channels_first',
#                       end_sparsity=0.,
#                       weight_decay=0.,
#                       name=None):
# def block_group(inputs,
#                 filters,
#                 block_fn,
#                 blocks,
#                 strides,
#                 is_training,
#                 name,
#                 pruning_method='baseline',
#                 init_method='baseline',
#                 data_format='channels_first',
#                 end_sparsity=0.,
#                 weight_decay=0.):
# def resnet_v1_generator(block_fn,
#                         num_blocks,
#                         num_classes,
#                         pruning_method='baseline',
#                         init_method='baseline',
#                         width=1.,
#                         prune_first_layer=True,
#                         prune_last_layer=True,
#                         data_format='channels_first',
#                         end_sparsity=0.,
#                         weight_decay=0.,
#                         name=None):
#   def model(inputs, is_training):
# def resnet_v1_(resnet_depth,
#                num_classes,
#                pruning_method='baseline',
#                init_method='baseline',
#                width=1.,
#                prune_first_layer=True,
#                prune_last_layer=True,
#                data_format='channels_first',
#                end_sparsity=0.,
#                weight_decay=0.,
#                name=None):
# class RandomSparseInitializer(init_ops.Initializer):
# class SparseConvVarianceScalingInitializer(init_ops.Initializer):
# class SparseFCVarianceScalingInitializer(init_ops.Initializer):


import functools
import tensorflow.compat.v1 as tf
from rigl.imagenet_resnet import resnet_model
from tensorflow.contrib import layers
      outputs. Useful to remove unnecessary dimensions for classification.
    name: Optional scope for the variables.
    global_pool: Optional boolean flag. If True, the input to the classification
      layer is avgpooled to size 1x1, for any input size. (This is not part
      of the original VGG architecture.)
    pruning_method: String that specifies the pruning method used to identify
      which weights to remove.
    init_method: ('baseline', 'sparse', 'random_zeros') Whether to use standard
      initialization or initialization that takes into the existing sparsity of
      the layer. 'sparse' only makes sense when combined with
      pruning_method == 'scratch'. 'random_zeros' set random weights to zero
      using end_sparsoty parameter and used with 'baseline' method.
    data_format: String specifying either "channels_first" for `[batch,
      channels, height, width]` or "channels_last for `[batch, height, width,
      channels]`.
    width: Float multiplier of the number of filters in each layer.
    prune_last_layer: Whether or not to prune the last layer.
    end_sparsity: Desired sparsity at the end of training. Necessary to
      initialize an already sparse network.
    weight_decay: Weight for the l2 regularization loss.

  Returns:
    net: the output of the logits layer (if num_classes is a non-zero integer),
      or the non-dropped-out input to the logits layer (if num_classes is 0 or
      None).
    end_points: a dict of tensors with intermediate activations. For
      backwards compatibility, some Tensors appear multiple times in the dict.
  """
  net_cfg = network_cfg[name]
  sparse_conv2d = functools.partial(

# Path: rigl/imagenet_resnet/resnet_model.py
# FLAGS = flags.FLAGS
# BATCH_NORM_DECAY = 0.9
# BATCH_NORM_EPSILON = 1e-5
# def batch_norm_relu(inputs, is_training, relu=True, init_zero=False,
#                     data_format='channels_first'):
# def fixed_padding(inputs, kernel_size, data_format='channels_first'):
#   def __init__(self, sparsity, seed=None, dtype=tf.float32):
#   def __call__(self, *args, **kwargs):
#   def get_config(self):
#   def __init__(self, sparsity, seed=None, dtype=tf.float32):
#   def __call__(self, shape, dtype=None, partition_info=None):
#   def get_config(self):
#   def __init__(self, sparsity, seed=None, dtype=tf.float32):
#   def __call__(self, shape, dtype=None, partition_info=None):
#   def get_config(self):
# def _pick_initializer(kernel_initializer, init_method, pruning_method,
#                       end_sparsity):
# def conv2d_fixed_padding(inputs,
#                          filters,
#                          kernel_size,
#                          strides,
#                          pruning_method='baseline',
#                          init_method='baseline',
#                          data_format='channels_first',
#                          end_sparsity=0.,
#                          weight_decay=0.,
#                          init_scale=1.0,
#                          name=None):
# def residual_block_(inputs,
#                     filters,
#                     is_training,
#                     strides,
#                     use_projection=False,
#                     pruning_method='baseline',
#                     init_method='baseline',
#                     data_format='channels_first',
#                     end_sparsity=0.,
#                     weight_decay=0.,
#                     name=''):
# def bottleneck_block_(inputs,
#                       filters,
#                       is_training,
#                       strides,
#                       use_projection=False,
#                       pruning_method='baseline',
#                       init_method='baseline',
#                       data_format='channels_first',
#                       end_sparsity=0.,
#                       weight_decay=0.,
#                       name=None):
# def block_group(inputs,
#                 filters,
#                 block_fn,
#                 blocks,
#                 strides,
#                 is_training,
#                 name,
#                 pruning_method='baseline',
#                 init_method='baseline',
#                 data_format='channels_first',
#                 end_sparsity=0.,
#                 weight_decay=0.):
# def resnet_v1_generator(block_fn,
#                         num_blocks,
#                         num_classes,
#                         pruning_method='baseline',
#                         init_method='baseline',
#                         width=1.,
#                         prune_first_layer=True,
#                         prune_last_layer=True,
#                         data_format='channels_first',
#                         end_sparsity=0.,
#                         weight_decay=0.,
#                         name=None):
#   def model(inputs, is_training):
# def resnet_v1_(resnet_depth,
#                num_classes,
#                pruning_method='baseline',
#                init_method='baseline',
#                width=1.,
#                prune_first_layer=True,
#                prune_last_layer=True,
#                data_format='channels_first',
#                end_sparsity=0.,
#                weight_decay=0.,
#                name=None):
# class RandomSparseInitializer(init_ops.Initializer):
# class SparseConvVarianceScalingInitializer(init_ops.Initializer):
# class SparseFCVarianceScalingInitializer(init_ops.Initializer):


# Path: rigl/imagenet_resnet/vgg.py
import functools
import tensorflow.compat.v1 as tf
from rigl.imagenet_resnet import resnet_model
from tensorflow.contrib import layers

      outputs. Useful to remove unnecessary dimensions for classification.
    name: Optional scope for the variables.
    global_pool: Optional boolean flag. If True, the input to the classification
      layer is avgpooled to size 1x1, for any input size. (This is not part
      of the original VGG architecture.)
    pruning_method: String that specifies the pruning method used to identify
      which weights to remove.
    init_method: ('baseline', 'sparse', 'random_zeros') Whether to use standard
      initialization or initialization that takes into the existing sparsity of
      the layer. 'sparse' only makes sense when combined with
      pruning_method == 'scratch'. 'random_zeros' set random weights to zero
      using end_sparsoty parameter and used with 'baseline' method.
    data_format: String specifying either "channels_first" for `[batch,
      channels, height, width]` or "channels_last for `[batch, height, width,
      channels]`.
    width: Float multiplier of the number of filters in each layer.
    prune_last_layer: Whether or not to prune the last layer.
    end_sparsity: Desired sparsity at the end of training. Necessary to
      initialize an already sparse network.
    weight_decay: Weight for the l2 regularization loss.

  Returns:
    net: the output of the logits layer (if num_classes is a non-zero integer),
      or the non-dropped-out input to the logits layer (if num_classes is 0 or
      None).
    end_points: a dict of tensors with intermediate activations. For
      backwards compatibility, some Tensors appear multiple times in the dict.
  """
  net_cfg = network_cfg[name]
  sparse_conv2d = functools.partial(

      resnet_model.conv2d_fixed_padding,