def train():
    """Training function."""
    trainer = gluon.Trainer(model.collect_params(), args.optimizer, {'learning_rate': args.lr})

    train_data_loader, val_data_loader, test_data_loader \
        = dataprocessor.make_dataloader(data_train, data_val, data_test, args)

    best_valid_bleu = 0.0
    for epoch_id in range(args.epochs):
        log_avg_loss = 0
        log_avg_gnorm = 0
        log_wc = 0
        log_start_time = time.time()
        for batch_id, (src_seq, tgt_seq, src_valid_length, tgt_valid_length)\
                in enumerate(train_data_loader):
            # logging.info(src_seq.context) Context suddenly becomes GPU.
            src_seq = src_seq.as_in_context(ctx)
            tgt_seq = tgt_seq.as_in_context(ctx)
            src_valid_length = src_valid_length.as_in_context(ctx)
            tgt_valid_length = tgt_valid_length.as_in_context(ctx)
            with mx.autograd.record():
                out, _ = model(src_seq, tgt_seq[:, :-1], src_valid_length, tgt_valid_length - 1)
                loss = loss_function(out, tgt_seq[:, 1:], tgt_valid_length - 1).mean()
                loss = loss * (tgt_seq.shape[1] - 1) / (tgt_valid_length - 1).mean()
                loss.backward()
            grads = [p.grad(ctx) for p in model.collect_params().values()]
            gnorm = gluon.utils.clip_global_norm(grads, args.clip)
            trainer.step(1)
            src_wc = src_valid_length.sum().asscalar()
            tgt_wc = (tgt_valid_length - 1).sum().asscalar()
            step_loss = loss.asscalar()
            log_avg_loss += step_loss
            log_avg_gnorm += gnorm
            log_wc += src_wc + tgt_wc
            if (batch_id + 1) % args.log_interval == 0:
                wps = log_wc / (time.time() - log_start_time)
                logging.info('[Epoch {} Batch {}/{}] loss={:.4f}, ppl={:.4f}, gnorm={:.4f}, '
                             'throughput={:.2f}K wps, wc={:.2f}K'
                             .format(epoch_id, batch_id + 1, len(train_data_loader),
                                     log_avg_loss / args.log_interval,
                                     np.exp(log_avg_loss / args.log_interval),
                                     log_avg_gnorm / args.log_interval,
                                     wps / 1000, log_wc / 1000))
                log_start_time = time.time()
                log_avg_loss = 0
                log_avg_gnorm = 0
                log_wc = 0
        valid_loss, valid_translation_out = evaluate(val_data_loader)
        valid_bleu_score, _, _, _, _ = compute_bleu([val_tgt_sentences], valid_translation_out)
        logging.info('[Epoch {}] valid Loss={:.4f}, valid ppl={:.4f}, valid bleu={:.2f}'
                     .format(epoch_id, valid_loss, np.exp(valid_loss), valid_bleu_score * 100))
        test_loss, test_translation_out = evaluate(test_data_loader)
        test_bleu_score, _, _, _, _ = compute_bleu([test_tgt_sentences], test_translation_out)
        logging.info('[Epoch {}] test Loss={:.4f}, test ppl={:.4f}, test bleu={:.2f}'
                     .format(epoch_id, test_loss, np.exp(test_loss), test_bleu_score * 100))
        dataprocessor.write_sentences(valid_translation_out,
                                      os.path.join(args.save_dir,
                                                   'epoch{:d}_valid_out.txt').format(epoch_id))
        dataprocessor.write_sentences(test_translation_out,
                                      os.path.join(args.save_dir,
                                                   'epoch{:d}_test_out.txt').format(epoch_id))
        if valid_bleu_score > best_valid_bleu:
            best_valid_bleu = valid_bleu_score
            save_path = os.path.join(args.save_dir, 'valid_best.params')
            logging.info('Save best parameters to {}'.format(save_path))
            model.save_parameters(save_path)
        if epoch_id + 1 >= (args.epochs * 2) // 3:
            new_lr = trainer.learning_rate * args.lr_update_factor
            logging.info('Learning rate change to {}'.format(new_lr))
            trainer.set_learning_rate(new_lr)
    if os.path.exists(os.path.join(args.save_dir, 'valid_best.params')):
        model.load_parameters(os.path.join(args.save_dir, 'valid_best.params'))
    valid_loss, valid_translation_out = evaluate(val_data_loader)
    valid_bleu_score, _, _, _, _ = compute_bleu([val_tgt_sentences], valid_translation_out)
    logging.info('Best model valid Loss={:.4f}, valid ppl={:.4f}, valid bleu={:.2f}'
                 .format(valid_loss, np.exp(valid_loss), valid_bleu_score * 100))
    test_loss, test_translation_out = evaluate(test_data_loader)
    test_bleu_score, _, _, _, _ = compute_bleu([test_tgt_sentences], test_translation_out)
    logging.info('Best model test Loss={:.4f}, test ppl={:.4f}, test bleu={:.2f}'
                 .format(test_loss, np.exp(test_loss), test_bleu_score * 100))
    dataprocessor.write_sentences(valid_translation_out,
                                  os.path.join(args.save_dir, 'best_valid_out.txt'))
    dataprocessor.write_sentences(test_translation_out,
                                  os.path.join(args.save_dir, 'best_test_out.txt'))
def train ( ) : trainer = gluon . Trainer ( model . collect_params ( ) , args . optimizer , { 'learning_rate' : args . lr } ) train_data_loader , val_data_loader , test_data_loader = dataprocessor . make_dataloader ( data_train , data_val , data_test , args ) best_valid_bleu = 0.0 for epoch_id in range ( args . epochs ) : log_avg_loss = 0 log_avg_gnorm = 0 log_wc = 0 log_start_time = time . time ( ) for batch_id , ( src_seq , tgt_seq , src_valid_length , tgt_valid_length ) in enumerate ( train_data_loader ) : src_seq = src_seq . as_in_context ( ctx ) tgt_seq = tgt_seq . as_in_context ( ctx ) src_valid_length = src_valid_length . as_in_context ( ctx ) tgt_valid_length = tgt_valid_length . as_in_context ( ctx ) with mx . autograd . record ( ) : out , _ = model ( src_seq , tgt_seq [ : , : - 1 ] , src_valid_length , tgt_valid_length - 1 ) loss = loss_function ( out , tgt_seq [ : , 1 : ] , tgt_valid_length - 1 ) . mean ( ) loss = loss * ( tgt_seq . shape [ 1 ] - 1 ) / ( tgt_valid_length - 1 ) . mean ( ) loss . backward ( ) grads = [ p . grad ( ctx ) for p in model . collect_params ( ) . values ( ) ] gnorm = gluon . utils . clip_global_norm ( grads , args . clip ) trainer . step ( 1 ) src_wc = src_valid_length . sum ( ) . asscalar ( ) tgt_wc = ( tgt_valid_length - 1 ) . sum ( ) . asscalar ( ) step_loss = loss . asscalar ( ) log_avg_loss += step_loss log_avg_gnorm += gnorm log_wc += src_wc + tgt_wc if ( batch_id + 1 ) % args . log_interval == 0 : wps = log_wc / ( time . time ( ) - log_start_time ) logging . info ( '[Epoch {} Batch {}/{}] loss={:.4f}, ppl={:.4f}, gnorm={:.4f}, ' 'throughput={:.2f}K wps, wc={:.2f}K' . format ( epoch_id , batch_id + 1 , len ( train_data_loader ) , log_avg_loss / args . log_interval , np . exp ( log_avg_loss / args . log_interval ) , log_avg_gnorm / args . log_interval , wps / 1000 , log_wc / 1000 ) ) log_start_time = time . time ( ) log_avg_loss = 0 log_avg_gnorm = 0 log_wc = 0 valid_loss , valid_translation_out = evaluate ( val_data_loader ) valid_bleu_score , _ , _ , _ , _ = compute_bleu ( [ val_tgt_sentences ] , valid_translation_out ) logging . info ( '[Epoch {}] valid Loss={:.4f}, valid ppl={:.4f}, valid bleu={:.2f}' . format ( epoch_id , valid_loss , np . exp ( valid_loss ) , valid_bleu_score * 100 ) ) test_loss , test_translation_out = evaluate ( test_data_loader ) test_bleu_score , _ , _ , _ , _ = compute_bleu ( [ test_tgt_sentences ] , test_translation_out ) logging . info ( '[Epoch {}] test Loss={:.4f}, test ppl={:.4f}, test bleu={:.2f}' . format ( epoch_id , test_loss , np . exp ( test_loss ) , test_bleu_score * 100 ) ) dataprocessor . write_sentences ( valid_translation_out , os . path . join ( args . save_dir , 'epoch{:d}_valid_out.txt' ) . format ( epoch_id ) ) dataprocessor . write_sentences ( test_translation_out , os . path . join ( args . save_dir , 'epoch{:d}_test_out.txt' ) . format ( epoch_id ) ) if valid_bleu_score > best_valid_bleu : best_valid_bleu = valid_bleu_score save_path = os . path . join ( args . save_dir , 'valid_best.params' ) logging . info ( 'Save best parameters to {}' . format ( save_path ) ) model . save_parameters ( save_path ) if epoch_id + 1 >= ( args . epochs * 2 ) // 3 : new_lr = trainer . learning_rate * args . lr_update_factor logging . info ( 'Learning rate change to {}' . format ( new_lr ) ) trainer . set_learning_rate ( new_lr ) if os . path . exists ( os . path . join ( args . save_dir , 'valid_best.params' ) ) : model . load_parameters ( os . path . join ( args . save_dir , 'valid_best.params' ) ) valid_loss , valid_translation_out = evaluate ( val_data_loader ) valid_bleu_score , _ , _ , _ , _ = compute_bleu ( [ val_tgt_sentences ] , valid_translation_out ) logging . info ( 'Best model valid Loss={:.4f}, valid ppl={:.4f}, valid bleu={:.2f}' . format ( valid_loss , np . exp ( valid_loss ) , valid_bleu_score * 100 ) ) test_loss , test_translation_out = evaluate ( test_data_loader ) test_bleu_score , _ , _ , _ , _ = compute_bleu ( [ test_tgt_sentences ] , test_translation_out ) logging . info ( 'Best model test Loss={:.4f}, test ppl={:.4f}, test bleu={:.2f}' . format ( test_loss , np . exp ( test_loss ) , test_bleu_score * 100 ) ) dataprocessor . write_sentences ( valid_translation_out , os . path . join ( args . save_dir , 'best_valid_out.txt' ) ) dataprocessor . write_sentences ( test_translation_out , os . path . join ( args . save_dir , 'best_test_out.txt' ) )
Training function.