def train():
    epochs = 100
    lr = 0.1
    lr_steps = [40, 70, np.inf]
    momentum = 0.9
    wd = 0.0005
    plot_period = 5
    ctx = [mx.gpu(i) for i in range(2)]
    batch_size = 256
    margin_s = 5
    margin_m = 0.2
    train_set = MNIST(train=True, transform=transform_train)
    train_data = gluon.data.DataLoader(train_set, batch_size, True, num_workers=4, last_batch='discard')
    val_set = MNIST(train=False, transform=transform_val)
    val_data = gluon.data.DataLoader(val_set, batch_size, shuffle=False, num_workers=4)
    net = LeNetPlus(embedding_size=64, feature_norm=True, weight_norm=True)
    net.initialize(init=mx.init.MSRAPrelu(), ctx=ctx)
    net.hybridize()
    loss = ArcLoss(s=margin_s, m=margin_m, classes=10)
    train_params = net.collect_params()
    trainer = gluon.Trainer(train_params, 'sgd', {'learning_rate': lr, 'momentum': momentum, 'wd': wd})
    lr_counter = 0
    metric = mtc.Accuracy()
    num_batch = len(train_data)
    for epoch in range((epochs + 1)):
        if (epoch == lr_steps[lr_counter]):
            trainer.set_learning_rate((trainer.learning_rate * 0.1))
            lr_counter += 1
        plot = False
        train_loss = 0
        metric.reset()
        tic = time.time()
        ebs = []
        lbs = []
        for batch in train_data:
            data = gluon.utils.split_and_load(batch[0], ctx_list=ctx, batch_axis=0, even_split=False)
            labels = gluon.utils.split_and_load(batch[1], ctx_list=ctx, batch_axis=0, even_split=False)
            with ag.record():
                ots = [net(X) for X in data]
                embedds = [ot[0] for ot in ots]
                outputs = [ot[1] for ot in ots]
                losses = [loss(yhat, y) for (yhat, y) in zip(outputs, labels)]
            for l in losses:
                ag.backward(l)
            if plot:
                for (es, ls) in zip(embedds, labels):
                    assert (len(es) == len(ls))
                    for idx in range(len(es)):
                        ebs.append(es[idx].asnumpy())
                        lbs.append(ls[idx].asscalar())
            trainer.step(batch_size)
            metric.update(labels, outputs)
            train_loss += (sum([l.mean().asscalar() for l in losses]) / len(losses))
        (_, train_acc) = metric.get()
        train_loss /= num_batch
        (val_acc, val_loss, val_ebs, val_lbs) = validate(net, val_data, ctx, loss, plot)
        if plot:
            ebs = np.vstack(ebs)
            lbs = np.hstack(lbs)
            plot_result(ebs, lbs, os.path.join('../../resources', 'arcloss-train-epoch{}.png'.format(epoch)))
            plot_result(val_ebs, val_lbs, os.path.join('../../resources', 'arcloss-val-epoch{}.png'.format(epoch)))
        toc = time.time()
        print(('[epoch % 3d] train accuracy: %.6f, train loss: %.6f | val accuracy: %.6f, val loss: %.6f, time: %.6f' % (epoch, train_acc, train_loss, val_acc, val_loss, (toc - tic))))
def train():
    epochs = 100
    lr = 0.1
    lr_steps = [40, 70, np.inf]
    momentum = 0.9
    wd = 0.0005
    plot_period = 5
    ctx = [mx.gpu(i) for i in range(2)]
    batch_size = 256
    margin_s = 5
    margin_m = 0.2
    train_set = MNIST(train=True, transform=transform_train)
    train_data = gluon.data.DataLoader(train_set, batch_size, True, num_workers=4, last_batch='discard')
    val_set = MNIST(train=False, transform=transform_val)
    val_data = gluon.data.DataLoader(val_set, batch_size, shuffle=False, num_workers=4)
    net = LeNetPlus(embedding_size=64, feature_norm=True, weight_norm=True)
    net.initialize(init=mx.init.MSRAPrelu(), ctx=ctx)
    net.hybridize()
    loss = ArcLoss(s=margin_s, m=margin_m, classes=10)
    train_params = net.collect_params()
    trainer = gluon.Trainer(train_params, 'sgd', {'learning_rate': lr, 'momentum': momentum, 'wd': wd})
    lr_counter = 0
    metric = mtc.Accuracy()
    num_batch = len(train_data)
    for epoch in range((epochs + 1)):
        if (epoch == lr_steps[lr_counter]):
            trainer.set_learning_rate((trainer.learning_rate * 0.1))
            lr_counter += 1
        plot = False
        train_loss = 0
        metric.reset()
        tic = time.time()
        ebs = []
        lbs = []
        for batch in train_data:
            data = gluon.utils.split_and_load(batch[0], ctx_list=ctx, batch_axis=0, even_split=False)
            labels = gluon.utils.split_and_load(batch[1], ctx_list=ctx, batch_axis=0, even_split=False)
            with ag.record():
                ots = [net(X) for X in data]
                embedds = [ot[0] for ot in ots]
                outputs = [ot[1] for ot in ots]
                losses = [loss(yhat, y) for (yhat, y) in zip(outputs, labels)]
            for l in losses:
                ag.backward(l)
            if plot:
                for (es, ls) in zip(embedds, labels):
                    assert (len(es) == len(ls))
                    for idx in range(len(es)):
                        ebs.append(es[idx].asnumpy())
                        lbs.append(ls[idx].asscalar())
            trainer.step(batch_size)
            metric.update(labels, outputs)
            train_loss += (sum([l.mean().asscalar() for l in losses]) / len(losses))
        (_, train_acc) = metric.get()
        train_loss /= num_batch
        (val_acc, val_loss, val_ebs, val_lbs) = validate(net, val_data, ctx, loss, plot)
        if plot:
            ebs = np.vstack(ebs)
            lbs = np.hstack(lbs)
            plot_result(ebs, lbs, os.path.join('../../resources', 'arcloss-train-epoch{}.png'.format(epoch)))
            plot_result(val_ebs, val_lbs, os.path.join('../../resources', 'arcloss-val-epoch{}.png'.format(epoch)))
        toc = time.time()
        print(('[epoch % 3d] train accuracy: %.6f, train loss: %.6f | val accuracy: %.6f, val loss: %.6f, time: %.6f' % (epoch, train_acc, train_loss, val_acc, val_loss, (toc - tic))))