def train():
ntasgd = False<EOL>best_val = float('<STR_LIT>')<EOL>start_train_time = time.time()<EOL>parameters = model.collect_params()<EOL>param_dict_avg = None<EOL>t = <NUM_LIT:0><EOL>avg_trigger = <NUM_LIT:0><EOL>n = <NUM_LIT:5><EOL>valid_losses = []<EOL>for epoch in range(args.epochs):<EOL><INDENT>total_L = <NUM_LIT:0.0><EOL>start_epoch_time = time.time()<EOL>start_log_interval_time = time.time()<EOL>hiddens = [model.begin_state(args.batch_size//len(context),<EOL>func=mx.nd.zeros, ctx=ctx) for ctx in context]<EOL>batch_i, i = <NUM_LIT:0>, <NUM_LIT:0><EOL>while i < len(train_data) - <NUM_LIT:1> - <NUM_LIT:1>:<EOL><INDENT>bptt = args.bptt if mx.nd.random.uniform().asscalar() < <NUM_LIT> else args.bptt / <NUM_LIT:2><EOL>seq_len = max(<NUM_LIT:5>, int(mx.nd.random.normal(bptt, <NUM_LIT:5>).asscalar()))<EOL>lr_batch_start = trainer.learning_rate<EOL>trainer.set_learning_rate(lr_batch_start*seq_len/args.bptt)<EOL>data, target = get_batch(train_data, i, seq_len=seq_len)<EOL>data_list = gluon.utils.split_and_load(data, context, batch_axis=<NUM_LIT:1>, even_split=True)<EOL>target_list = gluon.utils.split_and_load(target, context, batch_axis=<NUM_LIT:1>, even_split=True)<EOL>hiddens = detach(hiddens)<EOL>Ls = []<EOL>with autograd.record():<EOL><INDENT>for j, (X, y, h) in enumerate(zip(data_list, target_list, hiddens)):<EOL><INDENT>output, h, encoder_hs, dropped_encoder_hs = model(X, h)<EOL>l = joint_loss(output, y, encoder_hs, dropped_encoder_hs)<EOL>Ls.append(l / (len(context) * X.size))<EOL>hiddens[j] = h<EOL><DEDENT><DEDENT>for L in Ls:<EOL><INDENT>L.backward()<EOL><DEDENT>grads = [p.grad(d.context) for p in parameters.values() for d in data_list]<EOL>gluon.utils.clip_global_norm(grads, args.clip)<EOL>if args.ntasgd and ntasgd:<EOL><INDENT>if param_dict_avg is None:<EOL><INDENT>param_dict_avg = {k.split(model._prefix)[<NUM_LIT:1>]: v.data(context[<NUM_LIT:0>]).copy()<EOL>for k, v in parameters.items()}<EOL><DEDENT><DEDENT>trainer.step(<NUM_LIT:1>)<EOL>if args.ntasgd and ntasgd:<EOL><INDENT>gamma = <NUM_LIT:1.0> / max(<NUM_LIT:1>, epoch * (len(train_data) // args.bptt)<EOL>+ batch_i - avg_trigger + <NUM_LIT:2>)<EOL>for name, param_avg in param_dict_avg.items():<EOL><INDENT>param_avg[:] += gamma * (parameters['<STR_LIT>'.format(model._prefix, name)]<EOL>.data(context[<NUM_LIT:0>]) - param_avg)<EOL><DEDENT><DEDENT>total_L += sum([mx.nd.sum(L).asscalar() for L in Ls])<EOL>trainer.set_learning_rate(lr_batch_start)<EOL>if batch_i % args.log_interval == <NUM_LIT:0> and batch_i > <NUM_LIT:0>:<EOL><INDENT>cur_L = total_L / args.log_interval<EOL>print('<STR_LIT>'<EOL>'<STR_LIT>'<EOL>% (epoch, batch_i, len(train_data) // args.bptt, cur_L, math.exp(cur_L),<EOL>args.batch_size * args.log_interval<EOL>/ (time.time() - start_log_interval_time),<EOL>lr_batch_start * seq_len / args.bptt))<EOL>total_L = <NUM_LIT:0.0><EOL>start_log_interval_time = time.time()<EOL><DEDENT>i += seq_len<EOL>batch_i += <NUM_LIT:1><EOL><DEDENT>mx.nd.waitall()<EOL>print('<STR_LIT>' % (<EOL>epoch, (args.batch_size * len(train_data)) / (time.time() - start_epoch_time)))<EOL>if args.ntasgd and ntasgd:<EOL><INDENT>mx.nd.save('<STR_LIT>'.format(args.save), param_dict_avg)<EOL><DEDENT>else:<EOL><INDENT>model.save_parameters('<STR_LIT>'.format(args.save))<EOL><DEDENT>val_L = evaluate(val_data, val_batch_size, '<STR_LIT>'.format(args.save), context[<NUM_LIT:0>])<EOL>print('<STR_LIT>' % (<EOL>epoch, time.time() - start_epoch_time, val_L, math.exp(val_L),<EOL>trainer.learning_rate))<EOL>if args.ntasgd and avg_trigger == <NUM_LIT:0>:<EOL><INDENT>if t > n and val_L > min(valid_losses[-n:]):<EOL><INDENT>if param_dict_avg is None:<EOL><INDENT>param_dict_avg = {k.split(model._prefix)[<NUM_LIT:1>]: v.data(context[<NUM_LIT:0>]).copy()<EOL>for k, v in parameters.items()}<EOL><DEDENT>else:<EOL><INDENT>for k, v in parameters.items():<EOL><INDENT>param_dict_avg[k.split(model._prefix)[<NUM_LIT:1>]]= v.data(context[<NUM_LIT:0>]).copy()<EOL><DEDENT><DEDENT>avg_trigger = epoch * (len(train_data) // args.bptt) + len(train_data) // args.bptt<EOL>print('<STR_LIT>' % avg_trigger)<EOL>ntasgd = True<EOL><DEDENT>valid_losses.append(val_L)<EOL>t += <NUM_LIT:1><EOL><DEDENT>if val_L < best_val:<EOL><INDENT>update_lr_epoch = <NUM_LIT:0><EOL>best_val = val_L<EOL>if args.ntasgd and ntasgd:<EOL><INDENT>mx.nd.save(args.save, param_dict_avg)<EOL><DEDENT>else:<EOL><INDENT>model.save_parameters(args.save)<EOL><DEDENT>test_L = evaluate(test_data, test_batch_size, args.save, context[<NUM_LIT:0>])<EOL>print('<STR_LIT>'<EOL>% (epoch, test_L, math.exp(test_L)))<EOL><DEDENT>else:<EOL><INDENT>update_lr_epoch += <NUM_LIT:1><EOL>if update_lr_epoch % args.lr_update_interval == <NUM_LIT:0> and update_lr_epoch != <NUM_LIT:0>:<EOL><INDENT>lr_scale = trainer.learning_rate * args.lr_update_factor<EOL>print('<STR_LIT>' % lr_scale)<EOL>trainer.set_learning_rate(lr_scale)<EOL>update_lr_epoch = <NUM_LIT:0><EOL><DEDENT><DEDENT><DEDENT>print('<STR_LIT>'<EOL>% ((args.batch_size * len(train_data) * args.epochs) / (time.time() - start_train_time)))<EOL>
Training loop for awd language model.