text: # Copyright 2019 The TensorFlow Authors. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ==============================================================================
  """Full implementation of deep autoencoder experiment from original K-FAC paper.
  
  This script demonstrates training using KFAC optimizer, updating the damping
  parameter according to the Levenberg-Marquardt rule, and using the quadratic
  model method for adapting the learning rate and momentum parameters.
  """
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import math
  # Dependency imports
  from absl import flags
  import kfac
  import sonnet as snt
  import tensorflow.compat.v1 as tf
  
  from kfac.examples import mnist
  from kfac.python.ops.kfac_utils import data_reader
  from kfac.python.ops.kfac_utils import data_reader_alt
  
  
  # Model parameters
  _ENCODER_SIZES = [1000, 500, 250, 30]
  _DECODER_SIZES = [250, 500, 1000]
  _NONLINEARITY = tf.tanh  # Note: sigmoid cannot be used with the default init.
  _WEIGHTS_INITIALIZER = None  # Default init
  
  
  flags.DEFINE_integer('train_steps', 10000, 'Number of training steps.')
  
  flags.DEFINE_integer('inverse_update_period', 5,
                       '# of steps between computing inverse of Fisher factor '
                       'matrices.')
  
  flags.DEFINE_integer('cov_update_period', 1,
                       '# of steps between computing covaraiance matrices.')
  
  flags.DEFINE_integer('damping_adaptation_interval', 5,
                       '# of steps between updating the damping parameter.')
  
  flags.DEFINE_integer('num_burnin_steps', 5, 'Number of steps at the '
                       'start of training where the optimizer will only perform '
                       'cov updates.')
  
  flags.DEFINE_integer('seed', 12345, 'Random seed')
  
  flags.DEFINE_float('learning_rate', 3e-3,
                     'Learning rate to use when lrmu_adaptation="off".')
  
  flags.DEFINE_float('momentum', 0.9,
                     'Momentum decay value to use when '
                     'lrmu_adaptation="off" or "only_lr".')
  
  flags.DEFINE_float('damping', 1e-2, 'The fixed damping value to use. This is '
                     'ignored if adapt_damping is True.')
  
  flags.DEFINE_float('l2_reg', 1e-5,
                     'L2 regularization applied to weight matrices.')
  
  flags.DEFINE_boolean('update_damping_immediately', True, 'Adapt the damping '
                       'immediately after the parameter update (i.e. in the same '
                       'sess.run() call).  Only safe if everything is a resource '
                       'variable.')
  
  flags.DEFINE_boolean('use_batch_size_schedule', True,
                       'If True then we use the growing mini-batch schedule from '
                       'the original K-FAC paper.')
  
  flags.DEFINE_integer('batch_size', 1024,
                       'The size of the mini-batches to use if not using the '
                       'schedule.')
  
  flags.DEFINE_string('lrmu_adaptation', 'on',
                      'If set to "on" then we use the quadratic model '
                      'based learning-rate and momentum adaptation method from '
                      'the original paper. Note that this only works well in '
                      'practice when use_batch_size_schedule=True. Can also '
                      'be set to "off" and "only_lr", which turns '
                      'it off, or uses a version where the momentum parameter '
                      'is fixed (resp.).')
  
  flags.DEFINE_boolean('use_alt_data_reader', True,
                       'If True we use the alternative data reader for MNIST '
                       'that is faster for small datasets.')
  
  flags.DEFINE_string('device', '/gpu:0',
                      'The device to run the major ops on.')
  
  flags.DEFINE_boolean('adapt_damping', True,
                       'If True we use the LM rule for damping adaptation as '
                       'described in the original K-FAC paper.')
  
  # When using damping adaptation it is advisable to start with a high
  # value. This value is probably far too high to use for most neural nets
  # if you aren't using damping adaptation. (Although it always depends on
  # the scale of the loss.)
  flags.DEFINE_float('initial_damping', 150.0,
                     'The initial damping value to use when adapt_damping is '
                     'True.')
  
  flags.DEFINE_string('optimizer', 'kfac',
                      'The optimizer to use. Can be kfac or adam. If adam is '
                      'used the various K-FAC hyperparameter map roughly on to '
                      'their Adam equivalents.')
  
  flags.DEFINE_boolean('auto_register_layers', True,
                       'If True we use the automatic registration feature '
                       'which relies on scanning the TF graph. Otherwise '
                       'registration is done manually by this script during '
                       'the construction of the model.')
  
  flags.DEFINE_boolean('use_keras_model', False,
                       'If True, we use a Keras version of the autoencoder '
                       'model. Only works when auto_register_layers=True.')
  
  flags.DEFINE_boolean('use_sequential_for_keras', True,
                       'If True, we construct the Keras model using the '
                       'Sequential class.')
  
  flags.DEFINE_boolean('use_control_flow_v2', False, 'If True, we use Control '
                       'Flow V2. Defaults to False.')
  
  
  FLAGS = flags.FLAGS
  
  
  def make_train_op(minibatch,
                    batch_size,
                    batch_loss,
                    layer_collection,
                    loss_fn,
                    prev_train_batch=None,
                    placement_strategy=None,
                    print_logs=False,
                    tf_replicator=None):
    """Constructs optimizer and train op.
  
    Args:
      minibatch: A list/tuple of Tensors (typically representing the current
        mini-batch of input images and labels).
      batch_size: Tensor of shape (). Size of the training mini-batch.
      batch_loss: Tensor of shape (). Mini-batch loss tensor.
      layer_collection: LayerCollection object. Registry for model parameters.
        Required when using a K-FAC optimizer.
      loss_fn: Function which takes as input a mini-batch and returns the loss.
      prev_train_batch: `Tensor` of the previous training batch, can be accessed
        from the data_reader.CachedReader cached_batch property. (Default: None)
      placement_strategy: `str`, the placement_strategy argument for
        `KfacOptimizer`. (Default: None)
      print_logs: `Bool`. If True we print logs using K-FAC's built-in
        tf.print-based logs printer. (Default: False)
      tf_replicator: A Replicator object or None. If not None, K-FAC will set
          itself up to work inside of the provided TF-Replicator object.
          (Default: None)
  
    Returns:
      train_op: Op that can be used to update model parameters.
      optimizer: Optimizer used to produce train_op.
  
    Raises:
      ValueError: If layer_collection is None when K-FAC is selected as an
        optimization method.
    """
    global_step = tf.train.get_or_create_global_step()
  
    if FLAGS.optimizer == 'kfac':
      if FLAGS.lrmu_adaptation == 'on':
        learning_rate = None
        momentum = None
        momentum_type = 'qmodel'
      elif FLAGS.lrmu_adaptation == 'only_lr':
        learning_rate = None
        momentum = FLAGS.momentum
        momentum_type = 'qmodel_fixedmu'
      elif FLAGS.lrmu_adaptation == 'off':
        learning_rate = FLAGS.learning_rate
        momentum = FLAGS.momentum
        # momentum_type = 'regular'
        momentum_type = 'adam'
  
      if FLAGS.adapt_damping:
        damping = FLAGS.initial_damping
      else:
        damping = FLAGS.damping
  
      optimizer = kfac.PeriodicInvCovUpdateKfacOpt(
          invert_every=FLAGS.inverse_update_period,
          cov_update_every=FLAGS.cov_update_period,
          learning_rate=learning_rate,
          damping=damping,
          cov_ema_decay=0.95,
          momentum=momentum,
          momentum_type=momentum_type,
          layer_collection=layer_collection,
          batch_size=batch_size,
          num_burnin_steps=FLAGS.num_burnin_steps,
          adapt_damping=FLAGS.adapt_damping,
          l2_reg=FLAGS.l2_reg,
          placement_strategy=placement_strategy,
          print_logs=print_logs,
          tf_replicator=tf_replicator,
          # Note that many of the arguments below don't do anything when
          # adapt_damping=False.
          update_damping_immediately=FLAGS.update_damping_immediately,
          is_chief=True,
          prev_train_batch=prev_train_batch,  # We don't actually need this unless
                                              # update_damping_immediately is
                                              # False.
          loss=batch_loss,
          loss_fn=loss_fn,
          damping_adaptation_decay=0.95,
          damping_adaptation_interval=FLAGS.damping_adaptation_interval,
          min_damping=1e-6,
          train_batch=minibatch,
          )
  
    elif FLAGS.optimizer == 'adam':
      optimizer = tf.train.AdamOptimizer(
          learning_rate=FLAGS.learning_rate,
          beta1=FLAGS.momentum,
          epsilon=FLAGS.damping,
          beta2=0.99)
  
    return optimizer.minimize(batch_loss, global_step=global_step), optimizer
  
  
  class AutoEncoder(snt.AbstractModule):
    """Simple autoencoder module."""
  
    def __init__(self,
                 input_size,
                 regularizers=None,
                 initializers=None,
                 custom_getter=None,
                 name='AutoEncoder'):
      super(AutoEncoder, self).__init__(custom_getter=custom_getter, name=name)
  
      if initializers is None:
        initializers = {'w': tf.glorot_uniform_initializer(),
                        'b': tf.zeros_initializer()}
      if regularizers is None:
        regularizers = {'w': lambda w: FLAGS.l2_reg*tf.nn.l2_loss(w),
                        'b': lambda w: FLAGS.l2_reg*tf.nn.l2_loss(w),}
  
      with self._enter_variable_scope():
        self._encoder = snt.nets.MLP(
            output_sizes=_ENCODER_SIZES,
            regularizers=regularizers,
            initializers=initializers,
            custom_getter=custom_getter,
            activation=_NONLINEARITY,
            activate_final=False)
        self._decoder = snt.nets.MLP(
            output_sizes=_DECODER_SIZES + [input_size],
            regularizers=regularizers,
            initializers=initializers,
            custom_getter=custom_getter,
            activation=_NONLINEARITY,
            activate_final=False)
  
    def _build(self, inputs):
      code = self._encoder(inputs)
      output = self._decoder(code)
  
      return output
  
  
  class MLPManualReg(snt.AbstractModule):
  
    def __init__(self,
                 output_sizes,
                 regularizers=None,
                 initializers=None,
                 custom_getter=None,
                 activation=_NONLINEARITY,
                 activate_final=False,
                 name='MLP'):
  
      super(MLPManualReg, self).__init__(custom_getter=custom_getter, name=name)
  
      self._output_sizes = output_sizes
      self._activation = activation
      self._activate_final = activate_final
  
      with self._enter_variable_scope():
        self._layers = [snt.Linear(self._output_sizes[i],
                                   name='linear_{}'.format(i),
                                   initializers=initializers,
                                   regularizers=regularizers,
                                   custom_getter=custom_getter,
                                   use_bias=True)
                        for i in range(len(self._output_sizes))]
  
    def _build(self, inputs, layer_collection=None):
      net = inputs
      for i in range(len(self._output_sizes)):
        layer_inputs = net
        net = self._layers[i](net)
        layer_outputs = net
        params = (self._layers[i].w, self._layers[i].b)
  
        if layer_collection is not None:
          layer_collection.register_fully_connected(params,
                                                    layer_inputs,
                                                    layer_outputs,
                                                    reuse=False)
  
        if i < len(self._output_sizes) - 1 or self._activate_final:
          net = self._activation(net)
  
      return net
  
  
  class AutoEncoderManualReg(snt.AbstractModule):
    """Simple autoencoder module."""
  
    def __init__(self,
                 input_size,
                 regularizers=None,
                 initializers=None,
                 custom_getter=None,
                 name='AutoEncoder'):
      super(AutoEncoderManualReg, self).__init__(custom_getter=custom_getter,
                                                 name=name)
  
      if initializers is None:
        initializers = {'w': tf.glorot_uniform_initializer(),
                        'b': tf.zeros_initializer()}
      if regularizers is None:
        regularizers = {'w': lambda w: FLAGS.l2_reg*tf.nn.l2_loss(w),
                        'b': lambda w: FLAGS.l2_reg*tf.nn.l2_loss(w),}
  
      with self._enter_variable_scope():
        self._encoder = MLPManualReg(
            output_sizes=_ENCODER_SIZES,
            regularizers=regularizers,
            initializers=initializers,
            custom_getter=custom_getter,
            activation=_NONLINEARITY,
            activate_final=False)
        self._decoder = MLPManualReg(
            output_sizes=_DECODER_SIZES + [input_size],
            regularizers=regularizers,
            initializers=initializers,
            custom_getter=custom_getter,
            activation=_NONLINEARITY,
            activate_final=False)
  
    def _build(self, inputs, layer_collection=None):
      code = self._encoder(inputs, layer_collection=layer_collection)
      output = self._decoder(code, layer_collection=layer_collection)
  
      return output
  
  
  def get_keras_autoencoder(**input_kwargs):
    """Returns autoencoder made with Keras.
  
    Args:
      **input_kwargs: Arguments to pass to tf.keras.layers.Input. You must include
        either the 'shape' or 'tensor' kwarg.
  
    Returns:
      A tf.keras.Model, the Autoencoder.
    """
    layers = tf.keras.layers
    regularizers = tf.keras.regularizers
  
    dense_kwargs = {
        'kernel_initializer': tf.glorot_uniform_initializer(),
        'bias_initializer': tf.zeros_initializer(),
        'kernel_regularizer': regularizers.l2(l=FLAGS.l2_reg),
        'bias_regularizer': regularizers.l2(l=FLAGS.l2_reg),
    }
  
    if FLAGS.use_sequential_for_keras:
      model = tf.keras.Sequential()
      # Create Encoder
      model.add(layers.Input(**input_kwargs))
      for size in _ENCODER_SIZES[:-1]:
        model.add(layers.Dense(
            size, activation=_NONLINEARITY, **dense_kwargs))
      model.add(layers.Dense(_ENCODER_SIZES[-1], **dense_kwargs))
  
      # Create Decoder
      for size in _DECODER_SIZES:
        model.add(layers.Dense(size, activation=_NONLINEARITY, **dense_kwargs))
      model.add(layers.Dense(784, **dense_kwargs))
  
    else:
      # Make sure you always wrap the input in keras
      inputs = layers.Input(**input_kwargs)
  
      x = inputs
      # Create Encoder
      for size in _ENCODER_SIZES[:-1]:
        x = layers.Dense(size, activation=_NONLINEARITY, **dense_kwargs)(x)
      x = layers.Dense(_ENCODER_SIZES[-1], **dense_kwargs)(x)
  
      # Create Decoder
      for size in _DECODER_SIZES:
        x = layers.Dense(size, activation=_NONLINEARITY, **dense_kwargs)(x)
      x = layers.Dense(784, **dense_kwargs)(x)
  
      model = tf.keras.Model(inputs=inputs, outputs=x)
  
    return model
  
  
  def compute_squared_error(logits, targets):
    """Compute mean squared error."""
    return tf.reduce_sum(
        tf.reduce_mean(tf.square(targets - tf.nn.sigmoid(logits)), axis=0))
  
  
  def compute_loss(logits=None,
                   labels=None,
                   return_error=False,
                   model=None):
    """Compute loss value."""
    if FLAGS.use_keras_model:
      total_regularization_loss = tf.add_n(model.losses)
    else:
      graph_regularizers = tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES)
      total_regularization_loss = tf.add_n(graph_regularizers)
  
    loss_matrix = tf.nn.sigmoid_cross_entropy_with_logits(logits=logits,
                                                          labels=labels)
    loss = tf.reduce_sum(tf.reduce_mean(loss_matrix, axis=0))
    regularized_loss = loss + total_regularization_loss
  
    if return_error:
      squared_error = compute_squared_error(logits, labels)
      return regularized_loss, squared_error
  
    return regularized_loss
  
  
  def load_mnist():
    """Creates MNIST dataset and wraps it inside cached data reader.
  
    Returns:
      cached_reader: `data_reader.CachedReader` instance which wraps MNIST
        dataset.
      num_examples: int. The number of training examples.
    """
    # Wrap the data set into cached_reader which provides variable sized training
    # and caches the read train batch.
  
    if not FLAGS.use_alt_data_reader:
      # Version 1 using data_reader.py (slow!)
      dataset, num_examples = mnist.load_mnist_as_dataset(flatten_images=True)
      if FLAGS.use_batch_size_schedule:
        max_batch_size = num_examples
      else:
        max_batch_size = FLAGS.batch_size
  
      # Shuffle before repeat is correct unless you want repeat cases in the
      # same batch.
      dataset = (dataset.shuffle(num_examples).repeat()
                 .batch(max_batch_size).prefetch(5))
      dataset = tf.compat.v1.data.make_one_shot_iterator(dataset).get_next()
  
      # This version of CachedDataReader requires the dataset to be shuffled
      return data_reader.CachedDataReader(dataset, max_batch_size), num_examples
  
    else:
      # Version 2 using data_reader_alt.py (faster)
      images, labels, num_examples = mnist.load_mnist_as_tensors(
          flatten_images=True)
      dataset = (images, labels)
  
      # This version of CachedDataReader requires the dataset to NOT be shuffled
      return data_reader_alt.CachedDataReader(dataset, num_examples), num_examples
  
  
  def _get_batch_size_schedule(minibatch_maxsize):
    """Returns training batch size schedule."""
    minibatch_maxsize_targetiter = 500
    minibatch_startsize = 1000
  
    div = (float(minibatch_maxsize_targetiter-1)
           / math.log(float(minibatch_maxsize)/minibatch_startsize, 2))
    return [
        min(int(2.**(float(k)/div) * minibatch_startsize), minibatch_maxsize)
        for k in range(minibatch_maxsize_targetiter)
    ]
  
  
  def construct_train_quants():
    """Returns tensors and optimizer required to run the autoencoder."""
    with tf.device(FLAGS.device):
      # Load dataset.
      cached_reader, num_examples = load_mnist()
      batch_size_schedule = _get_batch_size_schedule(num_examples)
      batch_size = tf.placeholder(shape=(), dtype=tf.int32, name='batch_size')
  
      minibatch = cached_reader(batch_size)
      features, _ = minibatch
  
      if FLAGS.auto_register_layers:
        if FLAGS.use_keras_model:
          training_model = get_keras_autoencoder(tensor=features)
        else:
          training_model = AutoEncoder(784)
      else:
        training_model = AutoEncoderManualReg(784)
  
      layer_collection = kfac.LayerCollection()
  
      def loss_fn(minibatch, logits=None, return_error=False):
        features, _ = minibatch
        if logits is None:
          logits = training_model(features)
  
        return compute_loss(
            logits=logits,
            labels=features,
            return_error=return_error,
            model=training_model)
  
      if FLAGS.use_keras_model:
        logits = training_model.output
      else:
        if FLAGS.auto_register_layers:
          logits = training_model(features)
        else:
          logits = training_model(features, layer_collection=layer_collection)
  
      (batch_loss, batch_error) = loss_fn(
          minibatch, logits=logits, return_error=True)
  
      # Make sure never to confuse this with register_softmax_cross_entropy_loss!
      layer_collection.register_sigmoid_cross_entropy_loss(logits,
                                                           seed=FLAGS.seed + 1)
      if FLAGS.auto_register_layers:
        layer_collection.auto_register_layers()
  
      # Make training op
      train_op, opt = make_train_op(
          minibatch,
          batch_size,
          batch_loss,
          layer_collection,
          loss_fn=loss_fn,
          prev_train_batch=cached_reader.cached_batch)
  
    return train_op, opt, batch_loss, batch_error, batch_size_schedule, batch_size
  
  
  def main(_):
  
    # If using update_damping_immediately resource variables must be enabled.
    # Would recommend always enabling them anyway.
    if FLAGS.update_damping_immediately:
      tf.enable_resource_variables()
  
    if FLAGS.use_control_flow_v2:
      tf.enable_control_flow_v2()
  
    if not FLAGS.auto_register_layers and FLAGS.use_keras_model:
      raise ValueError('Require auto_register_layers=True when using Keras '
                       'model.')
  
    tf.set_random_seed(FLAGS.seed)
    (train_op, opt, batch_loss, batch_error, batch_size_schedule,
     batch_size) = construct_train_quants()
  
    global_step = tf.train.get_or_create_global_step()
  
    if FLAGS.optimizer == 'kfac':
      # We need to put the control depenency on train_op here so that we are
      # guaranteed to get the up-to-date values of these various quantities.
      # Otherwise there is a race condition and we might get the old values,
      # nondeterministically. Another solution would be to get these values in
      # a separate sess.run call, but this can sometimes cause problems with
      # training frameworks that use hooks (see the comments below).
      with tf.control_dependencies([train_op]):
        learning_rate = opt.learning_rate
        momentum = opt.momentum
        damping = opt.damping
        rho = opt.rho
        qmodel_change = opt.qmodel_change
  
    # Without setting allow_soft_placement=True there will be problems when
    # the optimizer tries to place certain ops like "mod" on the GPU (which isn't
    # supported).
    config = tf.ConfigProto(allow_soft_placement=True)
  
    # It's good practice to put everything into a single sess.run call. The
    # reason is that certain "training frameworks" like to run hooks at each
    # sess.run call, and there is an implicit expectation there will only
    # be one sess.run call every "iteration" of the "optimizer". For example,
    # a framework might try to print the loss at each sess.run call, causing
    # the mini-batch to be advanced, thus completely breaking the "cached
    # batch" mechanism that the damping adaptation method may rely on. (Plus
    # there will also be the extra cost of having to reevaluate the loss
    # twice.)  That being said we don't completely do that here because it's
    # inconvenient.
  
    # Train model.
    with tf.train.MonitoredTrainingSession(save_checkpoint_secs=30,
                                           config=config) as sess:
      for _ in range(FLAGS.train_steps):
        i = sess.run(global_step)
  
        if FLAGS.use_batch_size_schedule:
          batch_size_ = batch_size_schedule[min(i, len(batch_size_schedule) - 1)]
        else:
          batch_size_ = FLAGS.batch_size
  
        if FLAGS.optimizer == 'kfac':
          (_, batch_loss_, batch_error_, learning_rate_, momentum_, damping_,
           rho_, qmodel_change_) = sess.run([train_op, batch_loss, batch_error,
                                             learning_rate, momentum, damping,
                                             rho, qmodel_change],
                                            feed_dict={batch_size: batch_size_})
        else:
          _, batch_loss_, batch_error_ = sess.run(
              [train_op, batch_loss, batch_error],
              feed_dict={batch_size: batch_size_})
  
        # Print training stats.
        tf.logging.info(
            'iteration: %d', i)
        tf.logging.info(
            'mini-batch size: %d | mini-batch loss = %f | mini-batch error = %f ',
            batch_size_, batch_loss_, batch_error_)
  
        if FLAGS.optimizer == 'kfac':
          tf.logging.info(
              'learning_rate = %f | momentum = %f',
              learning_rate_, momentum_)
          tf.logging.info(
              'damping = %f | rho = %f | qmodel_change = %f',
              damping_, rho_, qmodel_change_)
  
        tf.logging.info('----')
  
  
  if __name__ == '__main__':
    tf.disable_v2_behavior()
    tf.app.run(main)
  