def train():
print(model)<EOL>from_epoch = <NUM_LIT:0><EOL>model.initialize(mx.init.Xavier(factor_type='<STR_LIT>'), ctx=context)<EOL>trainer_params = {'<STR_LIT>': args.lr, '<STR_LIT>': <NUM_LIT:0>, '<STR_LIT>': args.eps}<EOL>trainer = gluon.Trainer(model.collect_params(), '<STR_LIT>', trainer_params)<EOL>if args.from_epoch:<EOL><INDENT>from_epoch = args.from_epoch<EOL>checkpoint_name = '<STR_LIT>'%(args.save, format(from_epoch - <NUM_LIT:1>, '<STR_LIT>'))<EOL>model.load_parameters(checkpoint_name)<EOL>trainer.load_states('<STR_LIT>'%args.save)<EOL>print('<STR_LIT>'%(checkpoint_name))<EOL><DEDENT>model.hybridize(static_alloc=True, static_shape=True)<EOL>encoder_params = model.encoder.collect_params().values()<EOL>embedding_params = list(model.embedding.collect_params().values())<EOL>parallel_model = ParallelBigRNN(model, loss)<EOL>parallel = Parallel(len(context), parallel_model)<EOL>for epoch in range(from_epoch, args.epochs):<EOL><INDENT>sys.stdout.flush()<EOL>total_L = <NUM_LIT:0.0><EOL>start_epoch_time = time.time()<EOL>start_log_interval_time = time.time()<EOL>hiddens = [model.begin_state(batch_size=args.batch_size,<EOL>func=mx.nd.zeros, ctx=ctx) for ctx in context]<EOL>nbatch = <NUM_LIT:0><EOL>has_next = True<EOL>train_data_iter = iter(train_data)<EOL>data, target, mask, sample = next(train_data_iter)<EOL>while has_next:<EOL><INDENT>nbatch += <NUM_LIT:1><EOL>hiddens = detach(hiddens)<EOL>Ls = []<EOL>for _, batch in enumerate(zip(data, target, mask, sample, hiddens)):<EOL><INDENT>parallel.put(batch)<EOL><DEDENT>for _ in range(len(data)):<EOL><INDENT>hidden, ls = parallel.get()<EOL>index = context.index(hidden[<NUM_LIT:0>].context)<EOL>hiddens[index] = hidden<EOL>Ls.append(ls)<EOL><DEDENT>try:<EOL><INDENT>data, target, mask, sample = next(train_data_iter)<EOL><DEDENT>except StopIteration:<EOL><INDENT>has_next = False<EOL><DEDENT>for ctx in context:<EOL><INDENT>x = embedding_params[<NUM_LIT:0>].grad(ctx)<EOL>x[:] *= args.batch_size<EOL>encoder_grad = [p.grad(ctx) for p in encoder_params]<EOL>gluon.utils.clip_global_norm(encoder_grad, args.clip)<EOL><DEDENT>trainer.step(len(context))<EOL>total_L += sum([mx.nd.sum(L).asscalar() / args.bptt for L in Ls])<EOL>if nbatch % args.log_interval == <NUM_LIT:0>:<EOL><INDENT>cur_L = total_L / args.log_interval / len(context)<EOL>ppl = math.exp(cur_L) if cur_L < <NUM_LIT:100> else float('<STR_LIT>')<EOL>print('<STR_LIT>'<EOL>'<STR_LIT>'<EOL>%(epoch, nbatch, cur_L, ppl,<EOL>train_batch_size*args.log_interval/(time.time()-start_log_interval_time)))<EOL>total_L = <NUM_LIT:0.0><EOL>start_log_interval_time = time.time()<EOL>sys.stdout.flush()<EOL><DEDENT><DEDENT>end_epoch_time = time.time()<EOL>print('<STR_LIT>'%(epoch, end_epoch_time - start_epoch_time))<EOL>mx.nd.waitall()<EOL>checkpoint_name = '<STR_LIT>'%(args.save, format(epoch, '<STR_LIT>'))<EOL>model.save_parameters(checkpoint_name)<EOL>trainer.save_states('<STR_LIT>'%args.save)<EOL><DEDENT>
Training loop for language model.