text: # Licensed to the Apache Software Foundation (ASF) under one
  # or more contributor license agreements.  See the NOTICE file
  # distributed with this work for additional information
  # regarding copyright ownership.  The ASF licenses this file
  # to you 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.
  
  """ example train fit utility """
  import logging
  import os
  import time
  import re
  import math
  import mxnet as mx
  
  def get_epoch_size(args, kv):
      return math.ceil(int(args.num_examples / kv.num_workers) / args.batch_size)
  
  def _get_lr_scheduler(args, kv):
      if 'lr_factor' not in args or args.lr_factor >= 1:
          return (args.lr, None)
      epoch_size = get_epoch_size(args, kv)
      begin_epoch = args.load_epoch if args.load_epoch else 0
      if 'pow' in args.lr_step_epochs:
          lr = args.lr
          max_up = args.num_epochs * epoch_size
          pwr = float(re.sub('pow[- ]*', '', args.lr_step_epochs))
          poly_sched = mx.lr_scheduler.PolyScheduler(max_up, lr, pwr)
          return (lr, poly_sched)
      step_epochs = [int(l) for l in args.lr_step_epochs.split(',')]
      lr = args.lr
      for s in step_epochs:
          if begin_epoch >= s:
              lr *= args.lr_factor
      if lr != args.lr:
          logging.info('Adjust learning rate to %e for epoch %d',
                       lr, begin_epoch)
  
      steps = [epoch_size * (x - begin_epoch)
               for x in step_epochs if x - begin_epoch > 0]
      if steps:
          return (lr, mx.lr_scheduler.MultiFactorScheduler(step=steps, factor=args.lr_factor,
                                                           base_lr=args.lr))
      else:
          return (lr, None)
  
  def _load_model(args, rank=0):
      if 'load_epoch' not in args or args.load_epoch is None:
          return (None, None, None)
      assert args.model_prefix is not None
      model_prefix = args.model_prefix
      if rank > 0 and os.path.exists("%s-%d-symbol.json" % (model_prefix, rank)):
          model_prefix += "-%d" % (rank)
      sym, arg_params, aux_params = mx.model.load_checkpoint(
          model_prefix, args.load_epoch)
      logging.info('Loaded model %s_%04d.params', model_prefix, args.load_epoch)
      return (sym, arg_params, aux_params)
  
  
  def _save_model(args, rank=0):
      if args.model_prefix is None:
          return None
      return mx.callback.do_checkpoint(args.model_prefix if rank == 0 else "%s-%d" % (
          args.model_prefix, rank), period=args.save_period)
  
  
  def add_fit_args(parser):
      """
      parser : argparse.ArgumentParser
      return a parser added with args required by fit
      """
      train = parser.add_argument_group('Training', 'model training')
      train.add_argument('--network', type=str,
                         help='the neural network to use')
      train.add_argument('--num-layers', type=int,
                         help='number of layers in the neural network, \
                               required by some networks such as resnet')
      train.add_argument('--gpus', type=str,
                         help='list of gpus to run, e.g. 0 or 0,2,5. empty means using cpu')
      train.add_argument('--kv-store', type=str, default='device',
                         help='key-value store type')
      train.add_argument('--num-epochs', type=int, default=100,
                         help='max num of epochs')
      train.add_argument('--lr', type=float, default=0.1,
                         help='initial learning rate')
      train.add_argument('--lr-factor', type=float, default=0.1,
                         help='the ratio to reduce lr on each step')
      train.add_argument('--lr-step-epochs', type=str,
                         help='the epochs to reduce the lr, e.g. 30,60')
      train.add_argument('--initializer', type=str, default='default',
                         help='the initializer type')
      train.add_argument('--optimizer', type=str, default='sgd',
                         help='the optimizer type')
      train.add_argument('--mom', type=float, default=0.9,
                         help='momentum for sgd')
      train.add_argument('--wd', type=float, default=0.0001,
                         help='weight decay for sgd')
      train.add_argument('--batch-size', type=int, default=128,
                         help='the batch size')
      train.add_argument('--disp-batches', type=int, default=20,
                         help='show progress for every n batches')
      train.add_argument('--model-prefix', type=str,
                         help='model prefix')
      train.add_argument('--save-period', type=int, default=1, help='params saving period')
      parser.add_argument('--monitor', dest='monitor', type=int, default=0,
                          help='log network parameters every N iters if larger than 0')
      train.add_argument('--load-epoch', type=int,
                         help='load the model on an epoch using the model-load-prefix')
      train.add_argument('--top-k', type=int, default=0,
                         help='report the top-k accuracy. 0 means no report.')
      train.add_argument('--loss', type=str, default='',
                         help='show the cross-entropy or nll loss. ce strands for cross-entropy, nll-loss stands for likelihood loss')
      train.add_argument('--test-io', type=int, default=0,
                         help='1 means test reading speed without training')
      train.add_argument('--dtype', type=str, default='float32',
                         help='precision: float32 or float16')
      train.add_argument('--gc-type', type=str, default='none',
                         help='type of gradient compression to use, \
                               takes `2bit` or `none` for now')
      train.add_argument('--gc-threshold', type=float, default=0.5,
                         help='threshold for 2bit gradient compression')
      # additional parameters for large batch sgd
      train.add_argument('--macrobatch-size', type=int, default=0,
                         help='distributed effective batch size')
      train.add_argument('--warmup-epochs', type=int, default=5,
                         help='the epochs to ramp-up lr to scaled large-batch value')
      train.add_argument('--warmup-strategy', type=str, default='linear',
                         help='the ramping-up strategy for large batch sgd')
      train.add_argument('--profile-worker-suffix', type=str, default='',
                         help='profile workers actions into this file. During distributed training\
                               filename saved will be rank1_ followed by this suffix')
      train.add_argument('--profile-server-suffix', type=str, default='',
                         help='profile server actions into a file with name like rank1_ followed by this suffix \
                               during distributed training')
      return train
  
  
  def fit(args, network, data_loader, **kwargs):
      """
      train a model
      args : argparse returns
      network : the symbol definition of the nerual network
      data_loader : function that returns the train and val data iterators
      """
      # kvstore
      kv = mx.kvstore.create(args.kv_store)
      if args.gc_type != 'none':
          kv.set_gradient_compression({'type': args.gc_type,
                                       'threshold': args.gc_threshold})
      if args.profile_server_suffix:
          mx.profiler.set_config(filename=args.profile_server_suffix, profile_all=True, profile_process='server')
          mx.profiler.set_state(state='run', profile_process='server')
  
      if args.profile_worker_suffix:
          if kv.num_workers > 1:
              filename = 'rank' + str(kv.rank) + '_' + args.profile_worker_suffix
          else:
              filename = args.profile_worker_suffix
          mx.profiler.set_config(filename=filename, profile_all=True, profile_process='worker')
          mx.profiler.set_state(state='run', profile_process='worker')
  
      # logging
      head = '%(asctime)-15s Node[' + str(kv.rank) + '] %(message)s'
      logging.basicConfig(level=logging.DEBUG, format=head)
      logging.info('start with arguments %s', args)
      
      epoch_size = get_epoch_size(args, kv)
  
      # data iterators
      (train, val) = data_loader(args, kv)
      if 'dist' in args.kv_store and not 'async' in args.kv_store:
          logging.info('Resizing training data to %d batches per machine', epoch_size)
          # resize train iter to ensure each machine has same number of batches per epoch
          # if not, dist_sync can hang at the end with one machine waiting for other machines
          train = mx.io.ResizeIter(train, epoch_size)
  
      if args.test_io:
          tic = time.time()
          for i, batch in enumerate(train):
              if isinstance(batch, list):
                  for b in batch:
                      for j in b.data:
                          j.wait_to_read()
              else:
                  for j in batch.data:
                      j.wait_to_read()
              if (i + 1) % args.disp_batches == 0:
                  logging.info('Batch [%d]\tSpeed: %.2f samples/sec', i,
                               args.disp_batches * args.batch_size / (time.time() - tic))
                  tic = time.time()
          return
  
      # load model
      if 'arg_params' in kwargs and 'aux_params' in kwargs:
          arg_params = kwargs['arg_params']
          aux_params = kwargs['aux_params']
      else:
          sym, arg_params, aux_params = _load_model(args, kv.rank)
          if sym is not None:
              assert sym.tojson() == network.tojson()
  
      # save model
      checkpoint = _save_model(args, kv.rank)
  
      # devices for training
      devs = mx.cpu() if args.gpus is None or args.gpus == "" else [
          mx.gpu(int(i)) for i in args.gpus.split(',')]
  
      # learning rate
      lr, lr_scheduler = _get_lr_scheduler(args, kv)
  
      # create model
      model = mx.mod.Module(
          context=devs,
          symbol=network
      )
  
      lr_scheduler = lr_scheduler
      optimizer_params = {
          'learning_rate': lr,
          'wd': args.wd,
          'lr_scheduler': lr_scheduler,
          'multi_precision': True}
  
      # Only a limited number of optimizers have 'momentum' property
      has_momentum = {'sgd', 'dcasgd', 'nag', 'signum', 'lbsgd'}
      if args.optimizer in has_momentum:
          optimizer_params['momentum'] = args.mom
  
      monitor = mx.mon.Monitor(
          args.monitor, pattern=".*") if args.monitor > 0 else None
  
      # A limited number of optimizers have a warmup period
      has_warmup = {'lbsgd', 'lbnag'}
      if args.optimizer in has_warmup:
          nworkers = kv.num_workers
          if epoch_size < 1:
              epoch_size = 1
          macrobatch_size = args.macrobatch_size
          if macrobatch_size < args.batch_size * nworkers:
              macrobatch_size = args.batch_size * nworkers
          #batch_scale = round(float(macrobatch_size) / args.batch_size / nworkers +0.4999)
          batch_scale = math.ceil(
              float(macrobatch_size) / args.batch_size / nworkers)
          optimizer_params['updates_per_epoch'] = epoch_size
          optimizer_params['begin_epoch'] = args.load_epoch if args.load_epoch else 0
          optimizer_params['batch_scale'] = batch_scale
          optimizer_params['warmup_strategy'] = args.warmup_strategy
          optimizer_params['warmup_epochs'] = args.warmup_epochs
          optimizer_params['num_epochs'] = args.num_epochs
  
      if args.initializer == 'default':
          if args.network == 'alexnet':
              # AlexNet will not converge using Xavier
              initializer = mx.init.Normal()
              # VGG will not trend to converge using Xavier-Gaussian
          elif args.network and 'vgg' in args.network:
              initializer = mx.init.Xavier()
          else:
              initializer = mx.init.Xavier(
                  rnd_type='gaussian', factor_type="in", magnitude=2)
      # initializer   = mx.init.Xavier(factor_type="in", magnitude=2.34),
      elif args.initializer == 'xavier':
          initializer = mx.init.Xavier()
      elif args.initializer == 'msra':
          initializer = mx.init.MSRAPrelu()
      elif args.initializer == 'orthogonal':
          initializer = mx.init.Orthogonal()
      elif args.initializer == 'normal':
          initializer = mx.init.Normal()
      elif args.initializer == 'uniform':
          initializer = mx.init.Uniform()
      elif args.initializer == 'one':
          initializer = mx.init.One()
      elif args.initializer == 'zero':
          initializer = mx.init.Zero()
  
      # evaluation metrices
      eval_metrics = ['accuracy']
      if args.top_k > 0:
          eval_metrics.append(mx.metric.create(
              'top_k_accuracy', top_k=args.top_k))
  
      supported_loss = ['ce', 'nll_loss']
      if len(args.loss) > 0:
          # ce or nll loss is only applicable to softmax output
          loss_type_list = args.loss.split(',')
          if 'softmax_output' in network.list_outputs():
              for loss_type in loss_type_list:
                  loss_type = loss_type.strip()
                  if loss_type == 'nll':
                      loss_type = 'nll_loss'
                  if loss_type not in supported_loss:
                      logging.warning(loss_type + ' is not an valid loss type, only cross-entropy or ' \
                                      'negative likelihood loss is supported!')
                  else:
                      eval_metrics.append(mx.metric.create(loss_type))
          else:
              logging.warning("The output is not softmax_output, loss argument will be skipped!")
  
      # callbacks that run after each batch
      batch_end_callbacks = [mx.callback.Speedometer(
          args.batch_size, args.disp_batches)]
      if 'batch_end_callback' in kwargs:
          cbs = kwargs['batch_end_callback']
          batch_end_callbacks += cbs if isinstance(cbs, list) else [cbs]
  
      # run
      model.fit(train,
                begin_epoch=args.load_epoch if args.load_epoch else 0,
                num_epoch=args.num_epochs,
                eval_data=val,
                eval_metric=eval_metrics,
                kvstore=kv,
                optimizer=args.optimizer,
                optimizer_params=optimizer_params,
                initializer=initializer,
                arg_params=arg_params,
                aux_params=aux_params,
                batch_end_callback=batch_end_callbacks,
                epoch_end_callback=checkpoint,
                allow_missing=True,
                monitor=monitor)
  
      if args.profile_server_suffix:
          mx.profiler.set_state(state='run', profile_process='server')
      if args.profile_worker_suffix:
          mx.profiler.set_state(state='run', profile_process='worker')