def train_model(model, train_data_loader, val_data_loader, embedding, ctx, args):
    """
    Train model and validate/save every epoch.
    """
    logger.info(vars(args))

    # Initialization
    model.hybridize()
    model.collect_params().initialize(mx.init.Normal(0.01), ctx=ctx)
    model.word_emb.weight.set_data(embedding.idx_to_vec)
    # Fix word embedding
    model.word_emb.weight.grad_req = 'null'

    loss_func = gluon.loss.SoftmaxCrossEntropyLoss()
    trainer = gluon.Trainer(model.collect_params(), 'adagrad',
                            {'learning_rate': args.lr,
                             'wd': args.weight_decay,
                             'clip_gradient': 5})

    checkpoints_dir = os.path.join(args.output_dir, 'checkpoints')
    if not os.path.exists(checkpoints_dir):
        os.makedirs(checkpoints_dir)

    best_val_acc = 0.
    for epoch_id in range(args.epochs):
        avg_loss = 0.
        avg_acc = 0.
        for batch_id, example in enumerate(train_data_loader):
            s1, s2, label = example
            s1 = s1.as_in_context(ctx)
            s2 = s2.as_in_context(ctx)
            label = label.as_in_context(ctx)

            with autograd.record():
                output = model(s1, s2)
                loss = loss_func(output, label).mean()
            loss.backward()
            trainer.step(1)
            avg_loss += loss.sum().asscalar()

            pred = output.argmax(axis=1)
            acc = (pred == label.astype(np.float32)).mean()
            avg_acc += acc.asscalar()

            if (batch_id + 1) % args.print_interval == 0:
                avg_loss /= args.print_interval
                avg_acc /= args.print_interval
                logger.info('[Epoch {} Batch {}/{}] loss={:.4f}, acc={:.4f}'
                            .format(epoch_id, batch_id + 1, len(train_data_loader),
                                    avg_loss, avg_acc))
                avg_loss = 0.
                avg_acc = 0.

        # Validation
        val_loss, val_acc = test_model(model, val_data_loader, loss_func, ctx)
        if val_acc > best_val_acc:
            best_val_acc = val_acc
            checkpoint_path = os.path.join(args.output_dir, 'checkpoints', 'valid_best.params')
            model.save_parameters(checkpoint_path)
        logger.info('[Epoch {}] valid loss={:.4f}, valid acc={:.4f}, best valid acc={:.4f}'
                    .format(epoch_id, val_loss, val_acc, best_val_acc))

        # Save checkpoint of last epoch
        checkpoint_path = os.path.join(args.output_dir, 'checkpoints', 'last.params')
        model.save_parameters(checkpoint_path)
def train_model ( model , train_data_loader , val_data_loader , embedding , ctx , args ) : logger . info ( vars ( args ) ) model . hybridize ( ) model . collect_params ( ) . initialize ( mx . init . Normal ( 0.01 ) , ctx = ctx ) model . word_emb . weight . set_data ( embedding . idx_to_vec ) model . word_emb . weight . grad_req = 'null' loss_func = gluon . loss . SoftmaxCrossEntropyLoss ( ) trainer = gluon . Trainer ( model . collect_params ( ) , 'adagrad' , { 'learning_rate' : args . lr , 'wd' : args . weight_decay , 'clip_gradient' : 5 } ) checkpoints_dir = os . path . join ( args . output_dir , 'checkpoints' ) if not os . path . exists ( checkpoints_dir ) : os . makedirs ( checkpoints_dir ) best_val_acc = 0. for epoch_id in range ( args . epochs ) : avg_loss = 0. avg_acc = 0. for batch_id , example in enumerate ( train_data_loader ) : s1 , s2 , label = example s1 = s1 . as_in_context ( ctx ) s2 = s2 . as_in_context ( ctx ) label = label . as_in_context ( ctx ) with autograd . record ( ) : output = model ( s1 , s2 ) loss = loss_func ( output , label ) . mean ( ) loss . backward ( ) trainer . step ( 1 ) avg_loss += loss . sum ( ) . asscalar ( ) pred = output . argmax ( axis = 1 ) acc = ( pred == label . astype ( np . float32 ) ) . mean ( ) avg_acc += acc . asscalar ( ) if ( batch_id + 1 ) % args . print_interval == 0 : avg_loss /= args . print_interval avg_acc /= args . print_interval logger . info ( '[Epoch {} Batch {}/{}] loss={:.4f}, acc={:.4f}' . format ( epoch_id , batch_id + 1 , len ( train_data_loader ) , avg_loss , avg_acc ) ) avg_loss = 0. avg_acc = 0. val_loss , val_acc = test_model ( model , val_data_loader , loss_func , ctx ) if val_acc > best_val_acc : best_val_acc = val_acc checkpoint_path = os . path . join ( args . output_dir , 'checkpoints' , 'valid_best.params' ) model . save_parameters ( checkpoint_path ) logger . info ( '[Epoch {}] valid loss={:.4f}, valid acc={:.4f}, best valid acc={:.4f}' . format ( epoch_id , val_loss , val_acc , best_val_acc ) ) checkpoint_path = os . path . join ( args . output_dir , 'checkpoints' , 'last.params' ) model . save_parameters ( checkpoint_path )
Train model and validate/save every epoch.