Training function.

def train(epochs, ctx):
    """Training function."""
    if isinstance(ctx, mx.Context):
        ctx = [ctx]
    net.initialize(mx.init.Xavier(magnitude=2), ctx=ctx)

    opt_options = {'learning_rate': opt.lr, 'wd': opt.wd}
    if opt.optimizer == 'sgd':
        opt_options['momentum'] = 0.9
    if opt.optimizer == 'adam':
        opt_options['epsilon'] = 1e-7
    trainer = gluon.Trainer(net.collect_params(), opt.optimizer,
                            opt_options,
                            kvstore=opt.kvstore)
    if opt.lr_beta > 0.0:
        # Jointly train class-specific beta.
        # See "sampling matters in deep embedding learning" paper for details.
        beta.initialize(mx.init.Constant(opt.beta), ctx=ctx)
        trainer_beta = gluon.Trainer([beta], 'sgd',
                                     {'learning_rate': opt.lr_beta, 'momentum': 0.9},
                                     kvstore=opt.kvstore)

    loss = MarginLoss(margin=opt.margin, nu=opt.nu)

    best_val = 0.0
    for epoch in range(epochs):
        tic = time.time()
        prev_loss, cumulative_loss = 0.0, 0.0

        # Learning rate schedule.
        trainer.set_learning_rate(get_lr(opt.lr, epoch, steps, opt.factor))
        logging.info('Epoch %d learning rate=%f', epoch, trainer.learning_rate)
        if opt.lr_beta > 0.0:
            trainer_beta.set_learning_rate(get_lr(opt.lr_beta, epoch, steps, opt.factor))
            logging.info('Epoch %d beta learning rate=%f', epoch, trainer_beta.learning_rate)

        # Inner training loop.
        for i in range(200):
            batch = train_data.next()
            data = gluon.utils.split_and_load(batch.data[0], ctx_list=ctx, batch_axis=0)
            label = gluon.utils.split_and_load(batch.label[0], ctx_list=ctx, batch_axis=0)

            Ls = []
            with ag.record():
                for x, y in zip(data, label):
                    a_indices, anchors, positives, negatives, _ = net(x)

                    if opt.lr_beta > 0.0:
                        L = loss(anchors, positives, negatives, beta, y[a_indices])
                    else:
                        L = loss(anchors, positives, negatives, opt.beta, None)

                    # Store the loss and do backward after we have done forward
                    # on all GPUs for better speed on multiple GPUs.
                    Ls.append(L)
                    cumulative_loss += nd.mean(L).asscalar()

                for L in Ls:
                    L.backward()

            # Update.
            trainer.step(batch.data[0].shape[0])
            if opt.lr_beta > 0.0:
                trainer_beta.step(batch.data[0].shape[0])

            if (i+1) % opt.log_interval == 0:
                logging.info('[Epoch %d, Iter %d] training loss=%f' % (
                    epoch, i+1, cumulative_loss - prev_loss))
                prev_loss = cumulative_loss

        logging.info('[Epoch %d] training loss=%f'%(epoch, cumulative_loss))
        logging.info('[Epoch %d] time cost: %f'%(epoch, time.time()-tic))

        names, val_accs = test(ctx)
        for name, val_acc in zip(names, val_accs):
            logging.info('[Epoch %d] validation: %s=%f'%(epoch, name, val_acc))

        if val_accs[0] > best_val:
            best_val = val_accs[0]
            logging.info('Saving %s.' % opt.save_model_prefix)
            net.save_parameters('%s.params' % opt.save_model_prefix)
    return best_val