def train():
trainer = gluon.Trainer(model.collect_params(), args.optimizer,<EOL>{'<STR_LIT>': args.lr, '<STR_LIT>': <NUM_LIT>, '<STR_LIT>': <NUM_LIT>})<EOL>train_data_loader, val_data_loader, test_data_loader= dataprocessor.make_dataloader(data_train, data_val, data_test, args,<EOL>use_average_length=True, num_shards=len(ctx))<EOL>if args.bleu == '<STR_LIT>':<EOL><INDENT>bpe = bool(args.dataset != '<STR_LIT>' and args.dataset != '<STR_LIT>')<EOL>split_compound_word = bpe<EOL>tokenized = True<EOL><DEDENT>elif args.bleu == '<STR_LIT>' or args.bleu == '<STR_LIT>':<EOL><INDENT>bpe = False<EOL>split_compound_word = False<EOL>tokenized = False<EOL><DEDENT>else:<EOL><INDENT>raise NotImplementedError<EOL><DEDENT>best_valid_bleu = <NUM_LIT:0.0><EOL>step_num = <NUM_LIT:0><EOL>warmup_steps = args.warmup_steps<EOL>grad_interval = args.num_accumulated<EOL>model.collect_params().setattr('<STR_LIT>', '<STR_LIT>')<EOL>average_start = (len(train_data_loader) // grad_interval) * (args.epochs - args.average_start)<EOL>average_param_dict = None<EOL>model.collect_params().zero_grad()<EOL>parallel = Parallel(num_ctxs, parallel_model)<EOL>for epoch_id in range(args.epochs):<EOL><INDENT>log_avg_loss = <NUM_LIT:0><EOL>log_wc = <NUM_LIT:0><EOL>loss_denom = <NUM_LIT:0><EOL>step_loss = <NUM_LIT:0><EOL>log_start_time = time.time()<EOL>for batch_id, seqsin enumerate(train_data_loader):<EOL><INDENT>if batch_id % grad_interval == <NUM_LIT:0>:<EOL><INDENT>step_num += <NUM_LIT:1><EOL>new_lr = args.lr / math.sqrt(args.num_units)* min(<NUM_LIT:1.> / math.sqrt(step_num), step_num * warmup_steps ** (-<NUM_LIT>))<EOL>trainer.set_learning_rate(new_lr)<EOL><DEDENT>src_wc, tgt_wc, bs = np.sum([(shard[<NUM_LIT:2>].sum(), shard[<NUM_LIT:3>].sum(), shard[<NUM_LIT:0>].shape[<NUM_LIT:0>])<EOL>for shard in seqs], axis=<NUM_LIT:0>)<EOL>seqs = [[seq.as_in_context(context) for seq in shard]<EOL>for context, shard in zip(ctx, seqs)]<EOL>Ls = []<EOL>for seq in seqs:<EOL><INDENT>parallel.put((seq, args.batch_size))<EOL><DEDENT>Ls = [parallel.get() for _ in range(len(ctx))]<EOL>src_wc = src_wc.asscalar()<EOL>tgt_wc = tgt_wc.asscalar()<EOL>loss_denom += tgt_wc - bs<EOL>if batch_id % grad_interval == grad_interval - <NUM_LIT:1> orbatch_id == len(train_data_loader) - <NUM_LIT:1>:<EOL><INDENT>if average_param_dict is None:<EOL><INDENT>average_param_dict = {k: v.data(ctx[<NUM_LIT:0>]).copy() for k, v in<EOL>model.collect_params().items()}<EOL><DEDENT>trainer.step(float(loss_denom) / args.batch_size / <NUM_LIT>)<EOL>param_dict = model.collect_params()<EOL>param_dict.zero_grad()<EOL>if step_num > average_start:<EOL><INDENT>alpha = <NUM_LIT:1.> / max(<NUM_LIT:1>, step_num - average_start)<EOL>for name, average_param in average_param_dict.items():<EOL><INDENT>average_param[:] += alpha * (param_dict[name].data(ctx[<NUM_LIT:0>]) - average_param)<EOL><DEDENT><DEDENT><DEDENT>step_loss += sum([L.asscalar() for L in Ls])<EOL>if batch_id % grad_interval == grad_interval - <NUM_LIT:1> orbatch_id == len(train_data_loader) - <NUM_LIT:1>:<EOL><INDENT>log_avg_loss += step_loss / loss_denom * args.batch_size * <NUM_LIT><EOL>loss_denom = <NUM_LIT:0><EOL>step_loss = <NUM_LIT:0><EOL><DEDENT>log_wc += src_wc + tgt_wc<EOL>if (batch_id + <NUM_LIT:1>) % (args.log_interval * grad_interval) == <NUM_LIT:0>:<EOL><INDENT>wps = log_wc / (time.time() - log_start_time)<EOL>logging.info('<STR_LIT>'<EOL>'<STR_LIT>'<EOL>.format(epoch_id, batch_id + <NUM_LIT:1>, len(train_data_loader),<EOL>log_avg_loss / args.log_interval,<EOL>np.exp(log_avg_loss / args.log_interval),<EOL>wps / <NUM_LIT:1000>, log_wc / <NUM_LIT:1000>))<EOL>log_start_time = time.time()<EOL>log_avg_loss = <NUM_LIT:0><EOL>log_wc = <NUM_LIT:0><EOL><DEDENT><DEDENT>mx.nd.waitall()<EOL>valid_loss, valid_translation_out = evaluate(val_data_loader, ctx[<NUM_LIT:0>])<EOL>valid_bleu_score, _, _, _, _ = compute_bleu([val_tgt_sentences], valid_translation_out,<EOL>tokenized=tokenized, tokenizer=args.bleu,<EOL>split_compound_word=split_compound_word,<EOL>bpe=bpe)<EOL>logging.info('<STR_LIT>'<EOL>.format(epoch_id, valid_loss, np.exp(valid_loss), valid_bleu_score * <NUM_LIT:100>))<EOL>test_loss, test_translation_out = evaluate(test_data_loader, ctx[<NUM_LIT:0>])<EOL>test_bleu_score, _, _, _, _ = compute_bleu([test_tgt_sentences], test_translation_out,<EOL>tokenized=tokenized, tokenizer=args.bleu,<EOL>split_compound_word=split_compound_word,<EOL>bpe=bpe)<EOL>logging.info('<STR_LIT>'<EOL>.format(epoch_id, test_loss, np.exp(test_loss), test_bleu_score * <NUM_LIT:100>))<EOL>dataprocessor.write_sentences(valid_translation_out,<EOL>os.path.join(args.save_dir,<EOL>'<STR_LIT>').format(epoch_id))<EOL>dataprocessor.write_sentences(test_translation_out,<EOL>os.path.join(args.save_dir,<EOL>'<STR_LIT>').format(epoch_id))<EOL>if valid_bleu_score > best_valid_bleu:<EOL><INDENT>best_valid_bleu = valid_bleu_score<EOL>save_path = os.path.join(args.save_dir, '<STR_LIT>')<EOL>logging.info('<STR_LIT>'.format(save_path))<EOL>model.save_parameters(save_path)<EOL><DEDENT>save_path = os.path.join(args.save_dir, '<STR_LIT>'.format(epoch_id))<EOL>model.save_parameters(save_path)<EOL><DEDENT>save_path = os.path.join(args.save_dir, '<STR_LIT>')<EOL>mx.nd.save(save_path, average_param_dict)<EOL>if args.average_checkpoint:<EOL><INDENT>for j in range(args.num_averages):<EOL><INDENT>params = mx.nd.load(os.path.join(args.save_dir,<EOL>'<STR_LIT>'.format(args.epochs - j - <NUM_LIT:1>)))<EOL>alpha = <NUM_LIT:1.> / (j + <NUM_LIT:1>)<EOL>for k, v in model._collect_params_with_prefix().items():<EOL><INDENT>for c in ctx:<EOL><INDENT>v.data(c)[:] += alpha * (params[k].as_in_context(c) - v.data(c))<EOL><DEDENT><DEDENT><DEDENT>save_path = os.path.join(args.save_dir,<EOL>'<STR_LIT>'.format(args.num_averages))<EOL>model.save_parameters(save_path)<EOL><DEDENT>elif args.average_start > <NUM_LIT:0>:<EOL><INDENT>for k, v in model.collect_params().items():<EOL><INDENT>v.set_data(average_param_dict[k])<EOL><DEDENT>save_path = os.path.join(args.save_dir, '<STR_LIT>')<EOL>model.save_parameters(save_path)<EOL><DEDENT>else:<EOL><INDENT>model.load_parameters(os.path.join(args.save_dir, '<STR_LIT>'), ctx)<EOL><DEDENT>valid_loss, valid_translation_out = evaluate(val_data_loader, ctx[<NUM_LIT:0>])<EOL>valid_bleu_score, _, _, _, _ = compute_bleu([val_tgt_sentences], valid_translation_out,<EOL>tokenized=tokenized, tokenizer=args.bleu, bpe=bpe,<EOL>split_compound_word=split_compound_word)<EOL>logging.info('<STR_LIT>'<EOL>.format(valid_loss, np.exp(valid_loss), valid_bleu_score * <NUM_LIT:100>))<EOL>test_loss, test_translation_out = evaluate(test_data_loader, ctx[<NUM_LIT:0>])<EOL>test_bleu_score, _, _, _, _ = compute_bleu([test_tgt_sentences], test_translation_out,<EOL>tokenized=tokenized, tokenizer=args.bleu, bpe=bpe,<EOL>split_compound_word=split_compound_word)<EOL>logging.info('<STR_LIT>'<EOL>.format(test_loss, np.exp(test_loss), test_bleu_score * <NUM_LIT:100>))<EOL>dataprocessor.write_sentences(valid_translation_out,<EOL>os.path.join(args.save_dir, '<STR_LIT>'))<EOL>dataprocessor.write_sentences(test_translation_out,<EOL>os.path.join(args.save_dir, '<STR_LIT>'))<EOL>
Training function.