text: import re
  import sys
  sys.path.insert(0, "../../python")
  import time
  import logging
  import os.path
  
  import mxnet as mx
  import numpy as np
  from speechSGD import speechSGD
  from lstm_proj import lstm_unroll
  from io_util import BucketSentenceIter, TruncatedSentenceIter, DataReadStream
  from config_util import parse_args, get_checkpoint_path, parse_contexts
  
  
  # some constants
  METHOD_BUCKETING = 'bucketing'
  METHOD_TBPTT = 'truncated-bptt'
  
  def prepare_data(args):
      batch_size = args.config.getint('train', 'batch_size')
      num_hidden = args.config.getint('arch', 'num_hidden')
      num_hidden_proj = args.config.getint('arch', 'num_hidden_proj')
      num_lstm_layer = args.config.getint('arch', 'num_lstm_layer')
  
      init_c = [('l%d_init_c'%l, (batch_size, num_hidden)) for l in range(num_lstm_layer)]
      if num_hidden_proj > 0:
          init_h = [('l%d_init_h'%l, (batch_size, num_hidden_proj)) for l in range(num_lstm_layer)]
      else:
          init_h = [('l%d_init_h'%l, (batch_size, num_hidden)) for l in range(num_lstm_layer)]
  
      init_states = init_c + init_h
  
      file_train = args.config.get('data', 'train')
      file_dev = args.config.get('data', 'dev')
      file_format = args.config.get('data', 'format')
      feat_dim = args.config.getint('data', 'xdim')
  
      train_data_args = {
              "gpu_chunk": 32768,
              "lst_file": file_train,
              "file_format": file_format,
              "separate_lines": True
              }
  
      dev_data_args = {
              "gpu_chunk": 32768,
              "lst_file": file_dev,
              "file_format": file_format,
              "separate_lines": True
              }
  
      train_sets = DataReadStream(train_data_args, feat_dim)
      dev_sets = DataReadStream(dev_data_args, feat_dim)
  
      return (init_states, train_sets, dev_sets)
  
  def CrossEntropy(labels, preds):
      labels = labels.reshape((-1,))
      preds = preds.reshape((-1, preds.shape[1]))
      loss = 0.
      num_inst = 0
      for i in range(preds.shape[0]):
          label = labels[i]
  
          if label > 0:
              loss += -np.log(max(1e-10, preds[i][int(label)]))
              num_inst += 1
      return loss , num_inst
  
  def Acc_exclude_padding(labels, preds):
      labels = labels.reshape((-1,))
      preds = preds.reshape((-1, preds.shape[1]))
      sum_metric = 0
      num_inst = 0
      for i in range(preds.shape[0]):
          pred_label = np.argmax(preds[i], axis=0)
          label = labels[i]
  
          ind = np.nonzero(label.flat)
          pred_label_real = pred_label.flat[ind]
          label_real = label.flat[ind]
          sum_metric += (pred_label_real == label_real).sum()
          num_inst += len(pred_label_real)
      return sum_metric, num_inst
  
  class SimpleLRScheduler(mx.lr_scheduler.LRScheduler):
      """A simple lr schedule that simply return `dynamic_lr`. We will set `dynamic_lr`
      dynamically based on performance on the validation set.
      """
      def __init__(self, dynamic_lr, effective_sample_count=1, momentum=0.9, optimizer="sgd"):
          super(SimpleLRScheduler, self).__init__()
          self.dynamic_lr = dynamic_lr
          self.effective_sample_count = effective_sample_count
          self.momentum = momentum
          self.optimizer = optimizer
  
      def __call__(self, num_update):
          if self.optimizer == "speechSGD":
              return self.dynamic_lr / self.effective_sample_count, self.momentum
          else:
              return self.dynamic_lr / self.effective_sample_count
  
  def score_with_state_forwarding(module, eval_data, eval_metric):
      eval_data.reset()
      eval_metric.reset()
  
      for eval_batch in eval_data:
          module.forward(eval_batch, is_train=False)
          module.update_metric(eval_metric, eval_batch.label)
  
          # copy over states
          outputs = module.get_outputs()
          # outputs[0] is softmax, 1:end are states
          for i in range(1, len(outputs)):
              outputs[i].copyto(eval_data.init_state_arrays[i-1])
  
  
  def get_initializer(args):
      init_type = getattr(mx.initializer, args.config.get('train', 'initializer'))
      init_scale = args.config.getfloat('train', 'init_scale')
      if init_type is mx.initializer.Xavier:
          return mx.initializer.Xavier(magnitude=init_scale)
      return init_type(init_scale)
  
  
  def do_training(training_method, args, module, data_train, data_val):
      from distutils.dir_util import mkpath
      mkpath(os.path.dirname(get_checkpoint_path(args)))
  
      batch_size = data_train.batch_size
      batch_end_callbacks = [mx.callback.Speedometer(batch_size, 
                                                     args.config.getint('train', 'show_every'))]
      eval_allow_extra = True if training_method == METHOD_TBPTT else False
      eval_metric = [mx.metric.np(CrossEntropy, allow_extra_outputs=eval_allow_extra),
                     mx.metric.np(Acc_exclude_padding, allow_extra_outputs=eval_allow_extra)]
      eval_metric = mx.metric.create(eval_metric)
      optimizer = args.config.get('train', 'optimizer')
      momentum = args.config.getfloat('train', 'momentum')
      learning_rate = args.config.getfloat('train', 'learning_rate')
      lr_scheduler = SimpleLRScheduler(learning_rate, momentum=momentum, optimizer=optimizer)
  
      if training_method == METHOD_TBPTT:
          lr_scheduler.seq_len = data_train.truncate_len
  
      n_epoch = 0
      num_epoch = args.config.getint('train', 'num_epoch')
      learning_rate = args.config.getfloat('train', 'learning_rate')
      decay_factor = args.config.getfloat('train', 'decay_factor')
      decay_bound = args.config.getfloat('train', 'decay_lower_bound')
      clip_gradient = args.config.getfloat('train', 'clip_gradient')
      weight_decay = args.config.getfloat('train', 'weight_decay')
      if clip_gradient == 0:
          clip_gradient = None
  
      last_acc = -float("Inf")
      last_params = None
  
      module.bind(data_shapes=data_train.provide_data,
                  label_shapes=data_train.provide_label,
                  for_training=True)
      module.init_params(initializer=get_initializer(args))
  
      def reset_optimizer():
          if optimizer == "sgd" or optimizer == "speechSGD":
              module.init_optimizer(kvstore='device',
                                optimizer=args.config.get('train', 'optimizer'),
                                optimizer_params={'lr_scheduler': lr_scheduler,
                                                  'momentum': momentum,
                                                  'rescale_grad': 1.0,
                                                  'clip_gradient': clip_gradient,
                                                  'wd': weight_decay},
                                force_init=True)
          else:
              module.init_optimizer(kvstore='device',
                                optimizer=args.config.get('train', 'optimizer'),
                                optimizer_params={'lr_scheduler': lr_scheduler,
                                                  'rescale_grad': 1.0,
                                                  'clip_gradient': clip_gradient,
                                                  'wd': weight_decay},
                                force_init=True)
      reset_optimizer()
  
      while True:
          tic = time.time()
          eval_metric.reset()
  
          for nbatch, data_batch in enumerate(data_train):
              if training_method == METHOD_TBPTT:
                  lr_scheduler.effective_sample_count = data_train.batch_size * truncate_len
                  lr_scheduler.momentum = np.power(np.power(momentum, 1.0/(data_train.batch_size * truncate_len)), data_batch.effective_sample_count)
              else:
                  if data_batch.effective_sample_count is not None:
                      lr_scheduler.effective_sample_count = 1#data_batch.effective_sample_count
  
              module.forward_backward(data_batch)
              module.update()
              module.update_metric(eval_metric, data_batch.label)
  
              batch_end_params = mx.model.BatchEndParam(epoch=n_epoch, nbatch=nbatch,
                                                        eval_metric=eval_metric,
                                                        locals=None)
              for callback in batch_end_callbacks:
                  callback(batch_end_params)
  
              if training_method == METHOD_TBPTT:
                  # copy over states
                  outputs = module.get_outputs()
                  # outputs[0] is softmax, 1:end are states
                  for i in range(1, len(outputs)):
                      outputs[i].copyto(data_train.init_state_arrays[i-1])
  
          for name, val in eval_metric.get_name_value():
              logging.info('Epoch[%d] Train-%s=%f', n_epoch, name, val)
          toc = time.time()
          logging.info('Epoch[%d] Time cost=%.3f', n_epoch, toc-tic)
  
          data_train.reset()
  
          # test on eval data
          score_with_state_forwarding(module, data_val, eval_metric)
  
          # test whether we should decay learning rate
          curr_acc = None
          for name, val in eval_metric.get_name_value():
              logging.info("Epoch[%d] Dev-%s=%f", n_epoch, name, val)
              if name == 'CrossEntropy':
                  curr_acc = val
          assert curr_acc is not None, 'cannot find Acc_exclude_padding in eval metric'
  
          if n_epoch > 0 and lr_scheduler.dynamic_lr > decay_bound and curr_acc > last_acc:
              logging.info('Epoch[%d] !!! Dev set performance drops, reverting this epoch',
                           n_epoch)
              logging.info('Epoch[%d] !!! LR decay: %g => %g', n_epoch,
                           lr_scheduler.dynamic_lr, lr_scheduler.dynamic_lr / float(decay_factor))
  
              lr_scheduler.dynamic_lr /= decay_factor
              # we reset the optimizer because the internal states (e.g. momentum)
              # might already be exploded, so we want to start from fresh
              reset_optimizer()
              module.set_params(*last_params)
          else:
              last_params = module.get_params()
              last_acc = curr_acc
              n_epoch += 1
  
              # save checkpoints
              mx.model.save_checkpoint(get_checkpoint_path(args), n_epoch,
                                       module.symbol, *last_params)
  
          if n_epoch == num_epoch:
              break
  
  if __name__ == '__main__':
      args = parse_args()
      args.config.write(sys.stdout)
  
      training_method = args.config.get('train', 'method')
      contexts = parse_contexts(args)
  
      init_states, train_sets, dev_sets = prepare_data(args)
      state_names = [x[0] for x in init_states]
  
      batch_size = args.config.getint('train', 'batch_size')
      num_hidden = args.config.getint('arch', 'num_hidden')
      num_hidden_proj = args.config.getint('arch', 'num_hidden_proj')
      num_lstm_layer = args.config.getint('arch', 'num_lstm_layer')
      feat_dim = args.config.getint('data', 'xdim')
      label_dim = args.config.getint('data', 'ydim')
  
      logging.basicConfig(level=logging.DEBUG, format='%(asctime)-15s %(message)s')
  
      if training_method == METHOD_BUCKETING:
          buckets = args.config.get('train', 'buckets')
          buckets = list(map(int, re.split(r'\W+', buckets)))
          data_train = BucketSentenceIter(train_sets, buckets, batch_size, init_states, feat_dim=feat_dim)
          data_val   = BucketSentenceIter(dev_sets, buckets, batch_size, init_states, feat_dim=feat_dim)
  
          def sym_gen(seq_len):
              sym = lstm_unroll(num_lstm_layer, seq_len, feat_dim, num_hidden=num_hidden,
                                num_label=label_dim, num_hidden_proj=num_hidden_proj)
              data_names = ['data'] + state_names
              label_names = ['softmax_label']
              return (sym, data_names, label_names)
  
          module = mx.mod.BucketingModule(sym_gen,
                                          default_bucket_key=data_train.default_bucket_key,
                                          context=contexts)
          do_training(training_method, args, module, data_train, data_val)
      elif training_method == METHOD_TBPTT:
          truncate_len = args.config.getint('train', 'truncate_len')
          data_train = TruncatedSentenceIter(train_sets, batch_size, init_states,
                                             truncate_len=truncate_len, feat_dim=feat_dim)
          data_val = TruncatedSentenceIter(dev_sets, batch_size, init_states,
                                           truncate_len=truncate_len, feat_dim=feat_dim,
                                           do_shuffling=False, pad_zeros=True)
          sym = lstm_unroll(num_lstm_layer, truncate_len, feat_dim, num_hidden=num_hidden,
                            num_label=label_dim, output_states=True, num_hidden_proj=num_hidden_proj)
          data_names = [x[0] for x in data_train.provide_data]
          label_names = [x[0] for x in data_train.provide_label]
          module = mx.mod.Module(sym, context=contexts, data_names=data_names,
                                 label_names=label_names)
          do_training(training_method, args, module, data_train, data_val)
      else:
          raise RuntimeError('Unknown training method: %s' % training_method)
  
      print("="*80)
      print("Finished Training")
      print("="*80)
      args.config.write(sys.stdout)
  