content: """
  Keras implementation of CapsNet in Hinton's paper Dynamic Routing Between Capsules.
  The current version maybe only works for TensorFlow backend. Actually it will be straightforward to re-write to TF code.
  Adopting to other backends should be easy, but I have not tested this.
  
  Usage:
         python capsulenet-multi-gpu.py
         python capsulenet-multi-gpu.py --gpus 2
         ... ...
  
  Result:
      About 55 seconds per epoch on two GTX1080Ti GPU cards
  
  Author: Xifeng Guo, E-mail: `guoxifeng1990@163.com`, Github: `https://github.com/XifengGuo/CapsNet-Keras`
  """
  
  from keras import optimizers
  from keras import backend as K
  
  K.set_image_data_format('channels_last')
  
  from capsulenet import CapsNet, margin_loss, load_mnist, manipulate_latent, test
  
  
  def train(model, data, args):
      """
      Training a CapsuleNet
      :param model: the CapsuleNet model
      :param data: a tuple containing training and testing data, like `((x_train, y_train), (x_test, y_test))`
      :param args: arguments
      :return: The trained model
      """
      # unpacking the data
      (x_train, y_train), (x_test, y_test) = data
  
      # callbacks
      log = callbacks.CSVLogger(args.save_dir + '/log.csv')
      tb = callbacks.TensorBoard(log_dir=args.save_dir + '/tensorboard-logs',
                                 batch_size=args.batch_size, histogram_freq=args.debug)
      lr_decay = callbacks.LearningRateScheduler(schedule=lambda epoch: args.lr * (0.9 ** epoch))
  
      # compile the model
      model.compile(optimizer=optimizers.Adam(lr=args.lr),
                    loss=[margin_loss, 'mse'],
                    loss_weights=[1., args.lam_recon])
  
      """
      # Training without data augmentation:
      model.fit([x_train, y_train], [y_train, x_train], batch_size=args.batch_size, epochs=args.epochs,
                validation_data=[[x_test, y_test], [y_test, x_test]], callbacks=[log, tb, checkpoint, lr_decay])
      """
  
      # Begin: Training with data augmentation ---------------------------------------------------------------------#
      def train_generator(x, y, batch_size, shift_fraction=0.):
          train_datagen = ImageDataGenerator(width_shift_range=shift_fraction,
                                             height_shift_range=shift_fraction)  # shift up to 2 pixel for MNIST
          generator = train_datagen.flow(x, y, batch_size=batch_size)
          while 1:
              x_batch, y_batch = generator.next()
              yield ([x_batch, y_batch], [y_batch, x_batch])
  
      # Training with data augmentation. If shift_fraction=0., also no augmentation.
      model.fit_generator(generator=train_generator(x_train, y_train, args.batch_size, args.shift_fraction),
                          steps_per_epoch=int(y_train.shape[0] / args.batch_size),
                          epochs=args.epochs,
                          validation_data=[[x_test, y_test], [y_test, x_test]],
                          callbacks=[log, tb, lr_decay])
      # End: Training with data augmentation -----------------------------------------------------------------------#
  
      from utils import plot_log
      plot_log(args.save_dir + '/log.csv', show=True)
  
      return model
  
  
  if __name__ == "__main__":
      import numpy as np
      import tensorflow as tf
      import os
      from keras.preprocessing.image import ImageDataGenerator
      from keras import callbacks
      from keras.utils.vis_utils import plot_model
      from keras.utils import multi_gpu_model
  
      # setting the hyper parameters
      import argparse
      parser = argparse.ArgumentParser(description="Capsule Network on MNIST.")
      parser.add_argument('--epochs', default=50, type=int)
      parser.add_argument('--batch_size', default=300, type=int)
      parser.add_argument('--lam_recon', default=0.392, type=float,
                          help="The coefficient for the loss of decoder")
      parser.add_argument('-r', '--routings', default=3, type=int,
                          help="Number of iterations used in routing algorithm. should > 0")
      parser.add_argument('--shift_fraction', default=0.1, type=float,
                          help="Fraction of pixels to shift at most in each direction.")
      parser.add_argument('--debug', default=0, type=int,
                          help="Save weights by TensorBoard")
      parser.add_argument('--save_dir', default='./result')
      parser.add_argument('-t', '--testing', action='store_true',
                          help="Test the trained model on testing dataset")
      parser.add_argument('--digit', default=5, type=int,
                          help="Digit to manipulate")
      parser.add_argument('-w', '--weights', default=None,
                          help="The path of the saved weights. Should be specified when testing")
      parser.add_argument('--lr', default=0.001, type=float,
                          help="Initial learning rate")
      parser.add_argument('--gpus', default=2, type=int)
      args = parser.parse_args()
      print(args)
      if not os.path.exists(args.save_dir):
          os.makedirs(args.save_dir)
  
      # load data
      (x_train, y_train), (x_test, y_test) = load_mnist()
  
      # define model
      with tf.device('/cpu:0'):
          model, eval_model, manipulate_model = CapsNet(input_shape=x_train.shape[1:],
                                                        n_class=len(np.unique(np.argmax(y_train, 1))),
                                                        routings=args.routings)
      model.summary()
      plot_model(model, to_file=args.save_dir+'/model.png', show_shapes=True)
  
      # train or test
      if args.weights is not None:  # init the model weights with provided one
          model.load_weights(args.weights)
      if not args.testing:
          # define muti-gpu model
          multi_model = multi_gpu_model(model, gpus=args.gpus)
          train(model=multi_model, data=((x_train, y_train), (x_test, y_test)), args=args)
          model.save_weights(args.save_dir + '/trained_model.h5')
          print('Trained model saved to \'%s/trained_model.h5\'' % args.save_dir)
          test(model=eval_model, data=(x_test, y_test), args=args)
      else:  # as long as weights are given, will run testing
          if args.weights is None:
              print('No weights are provided. Will test using random initialized weights.')
          manipulate_latent(manipulate_model, (x_test, y_test), args)
          test(model=eval_model, data=(x_test, y_test), args=args)
  