text: #!/usr/bin/env python
  
  # Copyright 2017 Johns Hopkins University (Shinji Watanabe)
  #  Apache 2.0  (http://www.apache.org/licenses/LICENSE-2.0)
  
  
  import copy
  import json
  import logging
  import math
  import os
  
  # chainer related
  from chainer.datasets import TransformDataset
  from chainer import reporter as reporter_module
  from chainer import training
  from chainer.training import extensions
  
  # torch related
  import torch
  
  # espnet related
  from espnet.asr.asr_utils import adadelta_eps_decay
  from espnet.asr.asr_utils import add_results_to_json
  from espnet.asr.asr_utils import CompareValueTrigger
  from espnet.asr.asr_utils import get_model_conf
  from espnet.asr.asr_utils import make_batchset
  from espnet.asr.asr_utils import PlotAttentionReport
  from espnet.asr.asr_utils import restore_snapshot
  from espnet.asr.asr_utils import torch_load
  from espnet.asr.asr_utils import torch_resume
  from espnet.asr.asr_utils import torch_save
  from espnet.asr.asr_utils import torch_snapshot
  from espnet.nets.pytorch_backend.e2e_asr import E2E
  from espnet.nets.pytorch_backend.e2e_asr import pad_list
  from espnet.nets.pytorch_backend.e2e_asr import StreamingE2E
  
  from espnet.utils.training.iterators import ShufflingEnabler
  from espnet.utils.training.iterators import ToggleableShufflingMultiprocessIterator
  from espnet.utils.training.iterators import ToggleableShufflingSerialIterator
  
  from espnet.transform.transformation import using_transform_config
  from espnet.utils.io_utils import LoadInputsAndTargets
  
  # rnnlm
  import espnet.lm.pytorch_backend.extlm as extlm_pytorch
  import espnet.lm.pytorch_backend.lm as lm_pytorch
  
  # matplotlib related
  import matplotlib
  import numpy as np
  
  from espnet.utils.training.tensorboard_logger import TensorboardLogger
  from tensorboardX import SummaryWriter
  
  from espnet.utils.deterministic_utils import set_deterministic_pytorch
  from espnet.utils.training.train_utils import check_early_stop
  from espnet.utils.training.train_utils import set_early_stop
  
  matplotlib.use('Agg')
  
  REPORT_INTERVAL = 100
  
  
  class CustomEvaluator(extensions.Evaluator):
      """Custom Evaluator for Pytorch
  
      :param torch.nn.Module model : The model to evaluate
      :param chainer.dataset.Iterator : The train iterator
      :param target :
      :param CustomConverter converter : The batch converter
      :param torch.device device : The device used
      """
  
      def __init__(self, model, iterator, target, converter, device):
          super(CustomEvaluator, self).__init__(iterator, target)
          self.model = model
          self.converter = converter
          self.device = device
  
      # The core part of the update routine can be customized by overriding
      def evaluate(self):
          iterator = self._iterators['main']
  
          if self.eval_hook:
              self.eval_hook(self)
  
          if hasattr(iterator, 'reset'):
              iterator.reset()
              it = iterator
          else:
              it = copy.copy(iterator)
  
          summary = reporter_module.DictSummary()
  
          self.model.eval()
          with torch.no_grad():
              for batch in it:
                  observation = {}
                  with reporter_module.report_scope(observation):
                      # read scp files
                      # x: original json with loaded features
                      #    will be converted to chainer variable later
                      x = self.converter(batch, self.device)
                      self.model(*x)
                  summary.add(observation)
          self.model.train()
  
          return summary.compute_mean()
  
  
  class CustomUpdater(training.StandardUpdater):
      """Custom Updater for Pytorch
  
      :param torch.nn.Module model : The model to update
      :param int grad_clip_threshold : The gradient clipping value to use
      :param chainer.dataset.Iterator train_iter : The training iterator
      :param torch.optim.optimizer optimizer: The training optimizer
      :param CustomConverter converter: The batch converter
      :param torch.device device : The device to use
      :param int ngpu : The number of gpus to use
      """
  
      def __init__(self, model, grad_clip_threshold, train_iter,
                   optimizer, converter, device, ngpu):
          super(CustomUpdater, self).__init__(train_iter, optimizer)
          self.model = model
          self.grad_clip_threshold = grad_clip_threshold
          self.converter = converter
          self.device = device
          self.ngpu = ngpu
  
      # The core part of the update routine can be customized by overriding.
      def update_core(self):
          # When we pass one iterator and optimizer to StandardUpdater.__init__,
          # they are automatically named 'main'.
          train_iter = self.get_iterator('main')
          optimizer = self.get_optimizer('main')
  
          # Get the next batch ( a list of json files)
          batch = train_iter.next()
          x = self.converter(batch, self.device)
  
          # Compute the loss at this time step and accumulate it
          optimizer.zero_grad()  # Clear the parameter gradients
          loss = self.model(*x)[0]
          if self.ngpu > 1:
              loss = loss.sum() / self.ngpu
          loss.backward()  # Backprop
          loss.detach()  # Truncate the graph
          # compute the gradient norm to check if it is normal or not
          grad_norm = torch.nn.utils.clip_grad_norm_(
              self.model.parameters(), self.grad_clip_threshold)
          logging.info('grad norm={}'.format(grad_norm))
          if math.isnan(grad_norm):
              logging.warning('grad norm is nan. Do not update model.')
          else:
              optimizer.step()
  
  
  class CustomConverter(object):
      """Custom batch converter for Pytorch
  
      :param int subsampling_factor : The subsampling factor
      """
  
      def __init__(self, subsampling_factor=1, preprocess_conf=None):
          self.subsampling_factor = subsampling_factor
          self.load_inputs_and_targets = LoadInputsAndTargets(
              mode='asr', load_output=True, preprocess_conf=preprocess_conf)
          self.ignore_id = -1
  
      def transform(self, item):
          return self.load_inputs_and_targets(item)
  
      def __call__(self, batch, device):
          """Transforms a batch and send it to a device
  
          :param list batch: The batch to transform
          :param torch.device device: The device to send to
          :return: a tuple xs_pad, ilens, ys_pad
          :rtype (torch.Tensor, torch.Tensor, torch.Tensor)
          """
          # batch should be located in list
          assert len(batch) == 1
          xs, ys = batch[0]
  
          # perform subsampling
          if self.subsampling_factor > 1:
              xs = [x[::self.subsampling_factor, :] for x in xs]
  
          # get batch of lengths of input sequences
          ilens = np.array([x.shape[0] for x in xs])
  
          # perform padding and convert to tensor
          xs_pad = pad_list([torch.from_numpy(x).float() for x in xs], 0).to(device)
          ilens = torch.from_numpy(ilens).to(device)
          ys_pad = pad_list([torch.from_numpy(y).long() for y in ys], self.ignore_id).to(device)
  
          return xs_pad, ilens, ys_pad
  
  
  def train(args):
      """Train with the given args
  
      :param Namespace args: The program arguments
      """
      set_deterministic_pytorch(args)
  
      # check cuda availability
      if not torch.cuda.is_available():
          logging.warning('cuda is not available')
  
      # get input and output dimension info
      with open(args.valid_json, 'rb') as f:
          valid_json = json.load(f)['utts']
      utts = list(valid_json.keys())
      idim = int(valid_json[utts[0]]['input'][0]['shape'][1])
      odim = int(valid_json[utts[0]]['output'][0]['shape'][1])
      logging.info('#input dims : ' + str(idim))
      logging.info('#output dims: ' + str(odim))
  
      # specify attention, CTC, hybrid mode
      if args.mtlalpha == 1.0:
          mtl_mode = 'ctc'
          logging.info('Pure CTC mode')
      elif args.mtlalpha == 0.0:
          mtl_mode = 'att'
          logging.info('Pure attention mode')
      else:
          mtl_mode = 'mtl'
          logging.info('Multitask learning mode')
  
      # specify model architecture
      model = E2E(idim, odim, args)
      subsampling_factor = model.subsample[0]
  
      if args.rnnlm is not None:
          rnnlm_args = get_model_conf(args.rnnlm, args.rnnlm_conf)
          rnnlm = lm_pytorch.ClassifierWithState(
              lm_pytorch.RNNLM(
                  len(args.char_list), rnnlm_args.layer, rnnlm_args.unit))
          torch.load(args.rnnlm, rnnlm)
          model.rnnlm = rnnlm
  
      # write model config
      if not os.path.exists(args.outdir):
          os.makedirs(args.outdir)
      model_conf = args.outdir + '/model.json'
      with open(model_conf, 'wb') as f:
          logging.info('writing a model config file to ' + model_conf)
          f.write(json.dumps((idim, odim, vars(args)), indent=4, sort_keys=True).encode('utf_8'))
      for key in sorted(vars(args).keys()):
          logging.info('ARGS: ' + key + ': ' + str(vars(args)[key]))
  
      reporter = model.reporter
  
      # check the use of multi-gpu
      if args.ngpu > 1:
          model = torch.nn.DataParallel(model, device_ids=list(range(args.ngpu)))
          logging.info('batch size is automatically increased (%d -> %d)' % (
              args.batch_size, args.batch_size * args.ngpu))
          args.batch_size *= args.ngpu
  
      # set torch device
      device = torch.device("cuda" if args.ngpu > 0 else "cpu")
      model = model.to(device)
  
      # Setup an optimizer
      if args.opt == 'adadelta':
          optimizer = torch.optim.Adadelta(
              model.parameters(), rho=0.95, eps=args.eps,
              weight_decay=args.weight_decay)
      elif args.opt == 'adam':
          optimizer = torch.optim.Adam(model.parameters(),
                                       weight_decay=args.weight_decay)
  
      # FIXME: TOO DIRTY HACK
      setattr(optimizer, "target", reporter)
      setattr(optimizer, "serialize", lambda s: reporter.serialize(s))
  
      # Setup a converter
      converter = CustomConverter(subsampling_factor=subsampling_factor,
                                  preprocess_conf=args.preprocess_conf)
  
      # read json data
      with open(args.train_json, 'rb') as f:
          train_json = json.load(f)['utts']
      with open(args.valid_json, 'rb') as f:
          valid_json = json.load(f)['utts']
  
      use_sortagrad = args.sortagrad == -1 or args.sortagrad > 0
      # make minibatch list (variable length)
      train = make_batchset(train_json, args.batch_size,
                            args.maxlen_in, args.maxlen_out, args.minibatches,
                            min_batch_size=args.ngpu if args.ngpu > 1 else 1, shortest_first=use_sortagrad)
      valid = make_batchset(valid_json, args.batch_size,
                            args.maxlen_in, args.maxlen_out, args.minibatches,
                            min_batch_size=args.ngpu if args.ngpu > 1 else 1)
      # hack to make batchsize argument as 1
      # actual bathsize is included in a list
      if args.n_iter_processes > 0:
          train_iter = ToggleableShufflingMultiprocessIterator(
              TransformDataset(train, converter.transform),
              batch_size=1, n_processes=args.n_iter_processes, n_prefetch=8, maxtasksperchild=20,
              shuffle=not use_sortagrad)
          valid_iter = ToggleableShufflingMultiprocessIterator(
              TransformDataset(valid, converter.transform),
              batch_size=1, repeat=False, shuffle=False,
              n_processes=args.n_iter_processes, n_prefetch=8, maxtasksperchild=20)
      else:
          train_iter = ToggleableShufflingSerialIterator(
              TransformDataset(train, converter.transform),
              batch_size=1, shuffle=not use_sortagrad)
          valid_iter = ToggleableShufflingSerialIterator(
              TransformDataset(valid, converter.transform),
              batch_size=1, repeat=False, shuffle=False)
  
      # Set up a trainer
      updater = CustomUpdater(
          model, args.grad_clip, train_iter, optimizer, converter, device, args.ngpu)
      trainer = training.Trainer(
          updater, (args.epochs, 'epoch'), out=args.outdir)
  
      if use_sortagrad:
          trainer.extend(ShufflingEnabler([train_iter]),
                         trigger=(args.sortagrad if args.sortagrad != -1 else args.epochs, 'epoch'))
  
      # Resume from a snapshot
      if args.resume:
          logging.info('resumed from %s' % args.resume)
          torch_resume(args.resume, trainer)
  
      # Evaluate the model with the test dataset for each epoch
      trainer.extend(CustomEvaluator(model, valid_iter, reporter, converter, device))
  
      # Save attention weight each epoch
      if args.num_save_attention > 0 and args.mtlalpha != 1.0:
          data = sorted(list(valid_json.items())[:args.num_save_attention],
                        key=lambda x: int(x[1]['input'][0]['shape'][1]), reverse=True)
          if hasattr(model, "module"):
              att_vis_fn = model.module.calculate_all_attentions
          else:
              att_vis_fn = model.calculate_all_attentions
          att_reporter = PlotAttentionReport(
              att_vis_fn, data, args.outdir + "/att_ws",
              converter=converter, device=device)
          trainer.extend(att_reporter, trigger=(1, 'epoch'))
      else:
          att_reporter = None
  
      # Make a plot for training and validation values
      trainer.extend(extensions.PlotReport(['main/loss', 'validation/main/loss',
                                            'main/loss_ctc', 'validation/main/loss_ctc',
                                            'main/loss_att', 'validation/main/loss_att'],
                                           'epoch', file_name='loss.png'))
      trainer.extend(extensions.PlotReport(['main/acc', 'validation/main/acc'],
                                           'epoch', file_name='acc.png'))
      trainer.extend(extensions.PlotReport(['main/cer_ctc', 'validation/main/cer_ctc'],
                                           'epoch', file_name='cer.png'))
  
      # Save best models
      trainer.extend(extensions.snapshot_object(model, 'model.loss.best', savefun=torch_save),
                     trigger=training.triggers.MinValueTrigger('validation/main/loss'))
      if mtl_mode is not 'ctc':
          trainer.extend(extensions.snapshot_object(model, 'model.acc.best', savefun=torch_save),
                         trigger=training.triggers.MaxValueTrigger('validation/main/acc'))
  
      # save snapshot which contains model and optimizer states
      trainer.extend(torch_snapshot(), trigger=(1, 'epoch'))
  
      # epsilon decay in the optimizer
      if args.opt == 'adadelta':
          if args.criterion == 'acc' and mtl_mode is not 'ctc':
              trainer.extend(restore_snapshot(model, args.outdir + '/model.acc.best', load_fn=torch_load),
                             trigger=CompareValueTrigger(
                                 'validation/main/acc',
                                 lambda best_value, current_value: best_value > current_value))
              trainer.extend(adadelta_eps_decay(args.eps_decay),
                             trigger=CompareValueTrigger(
                                 'validation/main/acc',
                                 lambda best_value, current_value: best_value > current_value))
          elif args.criterion == 'loss':
              trainer.extend(restore_snapshot(model, args.outdir + '/model.loss.best', load_fn=torch_load),
                             trigger=CompareValueTrigger(
                                 'validation/main/loss',
                                 lambda best_value, current_value: best_value < current_value))
              trainer.extend(adadelta_eps_decay(args.eps_decay),
                             trigger=CompareValueTrigger(
                                 'validation/main/loss',
                                 lambda best_value, current_value: best_value < current_value))
  
      # Write a log of evaluation statistics for each epoch
      trainer.extend(extensions.LogReport(trigger=(REPORT_INTERVAL, 'iteration')))
      report_keys = ['epoch', 'iteration', 'main/loss', 'main/loss_ctc', 'main/loss_att',
                     'validation/main/loss', 'validation/main/loss_ctc', 'validation/main/loss_att',
                     'main/acc', 'validation/main/acc', 'main/cer_ctc', 'validation/main/cer_ctc',
                     'elapsed_time']
      if args.opt == 'adadelta':
          trainer.extend(extensions.observe_value(
              'eps', lambda trainer: trainer.updater.get_optimizer('main').param_groups[0]["eps"]),
              trigger=(REPORT_INTERVAL, 'iteration'))
          report_keys.append('eps')
      if args.report_cer:
          report_keys.append('validation/main/cer')
      if args.report_wer:
          report_keys.append('validation/main/wer')
      trainer.extend(extensions.PrintReport(
          report_keys), trigger=(REPORT_INTERVAL, 'iteration'))
  
      trainer.extend(extensions.ProgressBar(update_interval=REPORT_INTERVAL))
      set_early_stop(trainer, args)
  
      if args.tensorboard_dir is not None and args.tensorboard_dir != "":
          writer = SummaryWriter(log_dir=args.tensorboard_dir)
          trainer.extend(TensorboardLogger(writer, att_reporter))
      # Run the training
      trainer.run()
      check_early_stop(trainer, args.epochs)
  
  
  def recog(args):
      """Decode with the given args
  
      :param Namespace args: The program arguments
      """
      set_deterministic_pytorch(args)
      # read training config
      idim, odim, train_args = get_model_conf(args.model, args.model_conf)
  
      # load trained model parameters
      logging.info('reading model parameters from ' + args.model)
      model = E2E(idim, odim, train_args)
      torch_load(args.model, model)
      model.recog_args = args
  
      # read rnnlm
      if args.rnnlm:
          rnnlm_args = get_model_conf(args.rnnlm, args.rnnlm_conf)
          rnnlm = lm_pytorch.ClassifierWithState(
              lm_pytorch.RNNLM(
                  len(train_args.char_list), rnnlm_args.layer, rnnlm_args.unit))
          torch_load(args.rnnlm, rnnlm)
          rnnlm.eval()
      else:
          rnnlm = None
  
      if args.word_rnnlm:
          rnnlm_args = get_model_conf(args.word_rnnlm, args.word_rnnlm_conf)
          word_dict = rnnlm_args.char_list_dict
          char_dict = {x: i for i, x in enumerate(train_args.char_list)}
          word_rnnlm = lm_pytorch.ClassifierWithState(lm_pytorch.RNNLM(
              len(word_dict), rnnlm_args.layer, rnnlm_args.unit))
          torch_load(args.word_rnnlm, word_rnnlm)
          word_rnnlm.eval()
  
          if rnnlm is not None:
              rnnlm = lm_pytorch.ClassifierWithState(
                  extlm_pytorch.MultiLevelLM(word_rnnlm.predictor,
                                             rnnlm.predictor, word_dict, char_dict))
          else:
              rnnlm = lm_pytorch.ClassifierWithState(
                  extlm_pytorch.LookAheadWordLM(word_rnnlm.predictor,
                                                word_dict, char_dict))
  
      # gpu
      if args.ngpu == 1:
          gpu_id = range(args.ngpu)
          logging.info('gpu id: ' + str(gpu_id))
          model.cuda()
          if rnnlm:
              rnnlm.cuda()
  
      # read json data
      with open(args.recog_json, 'rb') as f:
          js = json.load(f)['utts']
      new_js = {}
  
      load_inputs_and_targets = LoadInputsAndTargets(
          mode='asr', load_output=False, sort_in_input_length=False,
          preprocess_conf=train_args.preprocess_conf
          if args.preprocess_conf is None else args.preprocess_conf)
  
      if args.batchsize == 0:
          with torch.no_grad():
              for idx, name in enumerate(js.keys(), 1):
                  logging.info('(%d/%d) decoding ' + name, idx, len(js.keys()))
                  batch = [(name, js[name])]
                  with using_transform_config({'train': True}):
                      feat = load_inputs_and_targets(batch)[0][0]
                  if args.streaming_window:
                      logging.info('Using streaming recognizer with window size %d frames', args.streaming_window)
                      se2e = StreamingE2E(e2e=model, recog_args=args, char_list=train_args.char_list, rnnlm=rnnlm)
                      for i in range(0, feat.shape[0], args.streaming_window):
                          logging.info('Feeding frames %d - %d', i, i + args.streaming_window)
                          se2e.accept_input(feat[i:i + args.streaming_window])
                      logging.info('Running offline attention decoder')
                      se2e.decode_with_attention_offline()
                      logging.info('Offline attention decoder finished')
                      nbest_hyps = se2e.retrieve_recognition()
                  else:
                      nbest_hyps = model.recognize(feat, args, train_args.char_list, rnnlm)
                  new_js[name] = add_results_to_json(js[name], nbest_hyps, train_args.char_list)
      else:
          try:
              from itertools import zip_longest as zip_longest
          except Exception:
              from itertools import izip_longest as zip_longest
  
          def grouper(n, iterable, fillvalue=None):
              kargs = [iter(iterable)] * n
              return zip_longest(*kargs, fillvalue=fillvalue)
  
          # sort data
          keys = list(js.keys())
          feat_lens = [js[key]['input'][0]['shape'][0] for key in keys]
          sorted_index = sorted(range(len(feat_lens)), key=lambda i: -feat_lens[i])
          keys = [keys[i] for i in sorted_index]
  
          with torch.no_grad():
              for names in grouper(args.batchsize, keys, None):
                  names = [name for name in names if name]
                  batch = [(name, js[name]) for name in names]
                  with using_transform_config({'train': False}):
                      feats = load_inputs_and_targets(batch)[0]
                  nbest_hyps = model.recognize_batch(feats, args, train_args.char_list, rnnlm=rnnlm)
                  for i, nbest_hyp in enumerate(nbest_hyps):
                      name = names[i]
                      new_js[name] = add_results_to_json(js[name], nbest_hyp, train_args.char_list)
  
      # TODO(watanabe) fix character coding problems when saving it
      with open(args.result_label, 'wb') as f:
          f.write(json.dumps({'utts': new_js}, indent=4, sort_keys=True).encode('utf_8'))
  