Training loop for language model.

def train():
    """Training loop for language model.
    """
    print(model)
    from_epoch = 0
    model.initialize(mx.init.Xavier(factor_type='out'), ctx=context)
    trainer_params = {'learning_rate': args.lr, 'wd': 0, 'eps': args.eps}
    trainer = gluon.Trainer(model.collect_params(), 'adagrad', trainer_params)
    if args.from_epoch:
        from_epoch = args.from_epoch
        checkpoint_name = '%s.%s'%(args.save, format(from_epoch - 1, '02d'))
        model.load_parameters(checkpoint_name)
        trainer.load_states('%s.state'%args.save)
        print('Loaded parameters from checkpoint %s'%(checkpoint_name))

    model.hybridize(static_alloc=True, static_shape=True)
    encoder_params = model.encoder.collect_params().values()
    embedding_params = list(model.embedding.collect_params().values())
    parallel_model = ParallelBigRNN(model, loss)
    parallel = Parallel(len(context), parallel_model)
    for epoch in range(from_epoch, args.epochs):
        sys.stdout.flush()
        total_L = 0.0
        start_epoch_time = time.time()
        start_log_interval_time = time.time()
        hiddens = [model.begin_state(batch_size=args.batch_size,
                                     func=mx.nd.zeros, ctx=ctx) for ctx in context]
        nbatch = 0
        has_next = True
        train_data_iter = iter(train_data)
        data, target, mask, sample = next(train_data_iter)

        while has_next:
            nbatch += 1
            hiddens = detach(hiddens)
            Ls = []
            for _, batch in enumerate(zip(data, target, mask, sample, hiddens)):
                parallel.put(batch)

            for _ in range(len(data)):
                hidden, ls = parallel.get()
                # hidden states are ordered by context id
                index = context.index(hidden[0].context)
                hiddens[index] = hidden
                Ls.append(ls)

            # prefetch the next batch of data
            try:
                data, target, mask, sample = next(train_data_iter)
            except StopIteration:
                has_next = False

            # rescale embedding grad
            for ctx in context:
                x = embedding_params[0].grad(ctx)
                x[:] *= args.batch_size
                encoder_grad = [p.grad(ctx) for p in encoder_params]
                # perform gradient clipping per ctx
                gluon.utils.clip_global_norm(encoder_grad, args.clip)

            trainer.step(len(context))

            total_L += sum([mx.nd.sum(L).asscalar() / args.bptt for L in Ls])

            if nbatch % args.log_interval == 0:
                cur_L = total_L / args.log_interval / len(context)
                ppl = math.exp(cur_L) if cur_L < 100 else float('inf')
                print('[Epoch %d Batch %d] loss %.2f, ppl %.2f, '
                      'throughput %.2f samples/s'
                      %(epoch, nbatch, cur_L, ppl,
                        train_batch_size*args.log_interval/(time.time()-start_log_interval_time)))
                total_L = 0.0
                start_log_interval_time = time.time()
                sys.stdout.flush()

        end_epoch_time = time.time()
        print('Epoch %d took %.2f seconds.'%(epoch, end_epoch_time - start_epoch_time))
        mx.nd.waitall()
        checkpoint_name = '%s.%s'%(args.save, format(epoch, '02d'))
        model.save_parameters(checkpoint_name)
        trainer.save_states('%s.state'%args.save)