# Buggy code
def main():
    opt = parse_args()

    filehandler = logging.FileHandler(opt.logging_file)
    streamhandler = logging.StreamHandler()

    logger = logging.getLogger("")
    logger.setLevel(logging.INFO)
    logger.addHandler(filehandler)
    logger.addHandler(streamhandler)

    logger.info(opt)

    bps.init()

    batch_size = opt.batch_size
    classes = 1000
    num_training_samples = 1281167

    num_gpus = opt.num_gpus
    # batch_size *= max(1, num_gpus)
    context = mx.gpu(bps.local_rank()) if num_gpus > 0 else mx.cpu(bps.local_rank())
    num_workers = opt.num_workers
    nworker = bps.size()
    rank = bps.rank()

    lr_decay = opt.lr_decay
    lr_decay_period = opt.lr_decay_period
    if opt.lr_decay_period > 0:
        lr_decay_epoch = list(range(lr_decay_period, opt.num_epochs, lr_decay_period))
    else:
        lr_decay_epoch = [int(i) for i in opt.lr_decay_epoch.split(",")]
    lr_decay_epoch = [e - opt.warmup_epochs for e in lr_decay_epoch]
    num_batches = num_training_samples // (batch_size * nworker)

    lr_scheduler = LRSequential(
        [
            LRScheduler(
                "linear",
                base_lr=opt.warmup_lr,
                target_lr=opt.lr * nworker / bps.local_size(),
                nepochs=opt.warmup_epochs,
                iters_per_epoch=num_batches,
            ),
            LRScheduler(
                opt.lr_mode,
                base_lr=opt.lr * nworker / bps.local_size(),
                target_lr=0,
                nepochs=opt.num_epochs - opt.warmup_epochs,
                iters_per_epoch=num_batches,
                step_epoch=lr_decay_epoch,
                step_factor=lr_decay,
                power=2,
            ),
        ]
    )

    model_name = opt.model

    kwargs = {"ctx": context, "pretrained": opt.use_pretrained, "classes": classes}
    if opt.use_gn:
        from gluoncv.nn import GroupNorm

        kwargs["norm_layer"] = GroupNorm
    if model_name.startswith("vgg"):
        kwargs["batch_norm"] = opt.batch_norm
    elif model_name.startswith("resnext"):
        kwargs["use_se"] = opt.use_se

    if opt.last_gamma:
        kwargs["last_gamma"] = True

    if opt.compressor:
        optimizer = "sgd"
    else:
        optimizer = "nag"

    optimizer_params = {
        "wd": opt.wd,
        "momentum": opt.momentum,
        "lr_scheduler": lr_scheduler,
    }

    if opt.dtype != "float32":
        optimizer_params["multi_precision"] = True

    net = get_model(model_name, **kwargs)
    net.cast(opt.dtype)
    if opt.resume_params is not "":
        net.load_parameters(opt.resume_params, ctx=context)

    # teacher model for distillation training
    if opt.teacher is not None and opt.hard_weight < 1.0:
        teacher_name = opt.teacher
        teacher = get_model(teacher_name, pretrained=True, classes=classes, ctx=context)
        teacher.cast(opt.dtype)
        distillation = True
    else:
        distillation = False

    # Two functions for reading data from record file or raw images
    def get_data_rec(
        rec_train, rec_train_idx, rec_val, rec_val_idx, batch_size, num_workers
    ):
        rec_train = os.path.expanduser(rec_train)
        rec_train_idx = os.path.expanduser(rec_train_idx)
        rec_val = os.path.expanduser(rec_val)
        rec_val_idx = os.path.expanduser(rec_val_idx)
        jitter_param = 0.4
        lighting_param = 0.1
        input_size = opt.input_size
        crop_ratio = opt.crop_ratio if opt.crop_ratio > 0 else 0.875
        resize = int(math.ceil(input_size / crop_ratio))
        mean_rgb = [123.68, 116.779, 103.939]
        std_rgb = [58.393, 57.12, 57.375]

        def batch_fn(batch, ctx):
            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
            )
            return data, label

        train_data = mx.io.ImageRecordIter(
            path_imgrec=rec_train,
            path_imgidx=rec_train_idx,
            preprocess_threads=num_workers,
            shuffle=True,
            batch_size=batch_size,
            data_shape=(3, input_size, input_size),
            mean_r=mean_rgb[0],
            mean_g=mean_rgb[1],
            mean_b=mean_rgb[2],
            std_r=std_rgb[0],
            std_g=std_rgb[1],
            std_b=std_rgb[2],
            rand_mirror=True,
            random_resized_crop=True,
            max_aspect_ratio=4.0 / 3.0,
            min_aspect_ratio=3.0 / 4.0,
            max_random_area=1,
            min_random_area=0.08,
            brightness=jitter_param,
            saturation=jitter_param,
            contrast=jitter_param,
            pca_noise=lighting_param,
            num_parts=nworker,
            part_index=rank,
        )
        val_data = mx.io.ImageRecordIter(
            path_imgrec=rec_val,
            path_imgidx=rec_val_idx,
            preprocess_threads=num_workers,
            shuffle=False,
            batch_size=batch_size,
            resize=resize,
            data_shape=(3, input_size, input_size),
            mean_r=mean_rgb[0],
            mean_g=mean_rgb[1],
            mean_b=mean_rgb[2],
            std_r=std_rgb[0],
            std_g=std_rgb[1],
            std_b=std_rgb[2],
            num_parts=nworker,
            part_index=rank,
        )
        return train_data, val_data, batch_fn

    def get_data_loader(data_dir, batch_size, num_workers):
        normalize = transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        jitter_param = 0.4
        lighting_param = 0.1
        input_size = opt.input_size
        crop_ratio = opt.crop_ratio if opt.crop_ratio > 0 else 0.875
        resize = int(math.ceil(input_size / crop_ratio))

        def batch_fn(batch, ctx):
            data = gluon.utils.split_and_load(batch[0], ctx_list=ctx, batch_axis=0)
            label = gluon.utils.split_and_load(batch[1], ctx_list=ctx, batch_axis=0)
            return data, label

        transform_train = transforms.Compose(
            [
                transforms.RandomResizedCrop(input_size),
                transforms.RandomFlipLeftRight(),
                transforms.RandomColorJitter(
                    brightness=jitter_param,
                    contrast=jitter_param,
                    saturation=jitter_param,
                ),
                transforms.RandomLighting(lighting_param),
                transforms.ToTensor(),
                normalize,
            ]
        )
        transform_test = transforms.Compose(
            [
                transforms.Resize(resize, keep_ratio=True),
                transforms.CenterCrop(input_size),
                transforms.ToTensor(),
                normalize,
            ]
        )

        train_data = gluon.data.DataLoader(
            imagenet.classification.ImageNet(data_dir, train=True).transform_first(
                transform_train
            ),
            batch_size=batch_size,
            shuffle=True,
            last_batch="discard",
            num_workers=num_workers,
        )
        val_data = gluon.data.DataLoader(
            imagenet.classification.ImageNet(data_dir, train=False).transform_first(
                transform_test
            ),
            batch_size=batch_size,
            shuffle=False,
            num_workers=num_workers,
        )

        return train_data, val_data, batch_fn

    if opt.use_rec:
        train_data, val_data, batch_fn = get_data_rec(
            opt.rec_train,
            opt.rec_train_idx,
            opt.rec_val,
            opt.rec_val_idx,
            batch_size,
            num_workers,
        )
    else:
        train_data, val_data, batch_fn = get_data_loader(
            opt.data_dir, batch_size, num_workers
        )

    if opt.mixup:
        train_metric = mx.metric.RMSE()
    else:
        train_metric = mx.metric.Accuracy()
    acc_top1 = mx.metric.Accuracy()
    acc_top5 = mx.metric.TopKAccuracy(5)

    save_frequency = opt.save_frequency
    if opt.save_dir and save_frequency:
        save_dir = opt.save_dir
        makedirs(save_dir)
    else:
        save_dir = ""
        save_frequency = 0

    def mixup_transform(label, classes, lam=1, eta=0.0):
        if isinstance(label, nd.NDArray):
            label = [label]
        res = []
        for l in label:
            y1 = l.one_hot(
                classes, on_value=1 - eta + eta / classes, off_value=eta / classes
            )
            y2 = l[::-1].one_hot(
                classes, on_value=1 - eta + eta / classes, off_value=eta / classes
            )
            res.append(lam * y1 + (1 - lam) * y2)
        return res

    def smooth(label, classes, eta=0.1):
        if isinstance(label, nd.NDArray):
            label = [label]
        smoothed = []
        for l in label:
            res = l.one_hot(
                classes, on_value=1 - eta + eta / classes, off_value=eta / classes
            )
            smoothed.append(res)
        return smoothed

    def test(ctx, val_data):
        if opt.use_rec:
            val_data.reset()
        acc_top1.reset()
        acc_top5.reset()
        for i, batch in enumerate(val_data):
            data, label = batch_fn(batch, ctx)
            outputs = [net(X.astype(opt.dtype, copy=False)) for X in data]
            acc_top1.update(label, outputs)
            acc_top5.update(label, outputs)

        _, top1 = acc_top1.get()
        _, top5 = acc_top5.get()
        return (1 - top1, 1 - top5)

    def train(ctx):
        if isinstance(ctx, mx.Context):
            ctx = [ctx]
        if opt.resume_params is "":
            net.initialize(mx.init.MSRAPrelu(), ctx=ctx)

        if opt.no_wd:
            for k, v in net.collect_params(".*beta|.*gamma|.*bias").items():
                v.wd_mult = 0.0

        compression_params = {
            "compressor": opt.compressor,
            "ef": opt.ef,
            "momentum": opt.compress_momentum,
            "scaling": opt.onebit_scaling,
        }

        trainer = bps.DistributedTrainer(
            net.collect_params(),
            optimizer,
            optimizer_params,
            compression_params=compression_params,
        )

        if opt.resume_states is not "":
            trainer.load_states(opt.resume_states)

        if opt.label_smoothing or opt.mixup:
            sparse_label_loss = False
        else:
            sparse_label_loss = True
        if distillation:
            L = gcv.loss.DistillationSoftmaxCrossEntropyLoss(
                temperature=opt.temperature,
                hard_weight=opt.hard_weight,
                sparse_label=sparse_label_loss,
            )
        else:
            L = gluon.loss.SoftmaxCrossEntropyLoss(sparse_label=sparse_label_loss)

        best_val_score = 1

        for epoch in range(opt.resume_epoch, opt.num_epochs):
            tic = time.time()
            if opt.use_rec:
                train_data.reset()
            train_metric.reset()
            btic = time.time()

            for i, batch in enumerate(train_data):
                data, label = batch_fn(batch, ctx)

                if opt.mixup:
                    lam = np.random.beta(opt.mixup_alpha, opt.mixup_alpha)
                    if epoch >= opt.num_epochs - opt.mixup_off_epoch:
                        lam = 1
                    data = [lam * X + (1 - lam) * X[::-1] for X in data]

                    if opt.label_smoothing:
                        eta = 0.1
                    else:
                        eta = 0.0
                    label = mixup_transform(label, classes, lam, eta)

                elif opt.label_smoothing:
                    hard_label = label
                    label = smooth(label, classes)

                if distillation:
                    teacher_prob = [
                        nd.softmax(
                            teacher(X.astype(opt.dtype, copy=False)) / opt.temperature
                        )
                        for X in data
                    ]

                with ag.record():
                    outputs = [net(X.astype(opt.dtype, copy=False)) for X in data]
                    if distillation:
                        loss = [
                            L(
                                yhat.astype("float32", copy=False),
                                y.astype("float32", copy=False),
                                p.astype("float32", copy=False),
                            )
                            for yhat, y, p in zip(outputs, label, teacher_prob)
                        ]
                    else:
                        loss = [
                            L(yhat, y.astype(opt.dtype, copy=False))
                            for yhat, y in zip(outputs, label)
                        ]
                for l in loss:
                    l.backward()
                trainer.step(batch_size)

                if opt.mixup:
                    output_softmax = [
                        nd.SoftmaxActivation(out.astype("float32", copy=False))
                        for out in outputs
                    ]
                    train_metric.update(label, output_softmax)
                else:
                    if opt.label_smoothing:
                        train_metric.update(hard_label, outputs)
                    else:
                        train_metric.update(label, outputs)

                if opt.log_interval and not (i + 1) % opt.log_interval:
                    train_metric_name, train_metric_score = train_metric.get()
                    logger.info(
                        "Epoch[%d] Batch [%d]\tSpeed: %f samples/sec\t%s=%f\tlr=%f\ttime=%f"
                        % (
                            epoch,
                            i,
                            batch_size
                            * nworker
                            * opt.log_interval
                            / (time.time() - btic),
                            train_metric_name,
                            train_metric_score,
                            trainer.learning_rate,
                            time.time() - btic,
                        )
                    )
                    btic = time.time()

            train_metric_name, train_metric_score = train_metric.get()
            throughput = int(batch_size * nworker * i / (time.time() - tic))

            logger.info(
                "[Epoch %d] training: %s=%f"
                % (epoch, train_metric_name, train_metric_score)
            )
            logger.info(
                "[Epoch %d] speed: %d samples/sec\ttime cost: %f"
                % (epoch, throughput, time.time() - tic)
            )

            err_top1_val, err_top5_val = test(ctx, val_data)

            logger.info(
                "[Epoch %d] validation: err-top1=%f err-top5=%f"
                % (epoch, err_top1_val, err_top5_val)
            )

            if err_top1_val < best_val_score:
                best_val_score = err_top1_val
                net.save_parameters(
                    "%s/%.4f-imagenet-%s-%d-best.params"
                    % (save_dir, best_val_score, model_name, epoch)
                )
                trainer.save_states(
                    "%s/%.4f-imagenet-%s-%d-best.states"
                    % (save_dir, best_val_score, model_name, epoch)
                )

            if save_frequency and save_dir and (epoch + 1) % save_frequency == 0:
                net.save_parameters(
                    "%s/imagenet-%s-%d.params" % (save_dir, model_name, epoch)
                )
                trainer.save_states(
                    "%s/imagenet-%s-%d.states" % (save_dir, model_name, epoch)
                )

        if save_frequency and save_dir:
            net.save_parameters(
                "%s/imagenet-%s-%d.params" % (save_dir, model_name, opt.num_epochs - 1)
            )
            trainer.save_states(
                "%s/imagenet-%s-%d.states" % (save_dir, model_name, opt.num_epochs - 1)
            )

    if opt.mode == "hybrid":
        net.hybridize(static_alloc=True, static_shape=True)
        if distillation:
            teacher.hybridize(static_alloc=True, static_shape=True)
    train(context)


# Traceback (TypeError)
In [1]: import tensorflow as tf

In [2]: import byteps.tensorflow as bps
WARNING: Logging before flag parsing goes to stderr.
W0627 11:36:47.010180 139917697820480 deprecation_wrapper.py:119] From /private/home/yuxinwu/.local/lib/python3.6/site-packages/byteps/tensorflow/__init__.py:79: The name tf.train.SessionRunHook is deprecated. Please use tf.estimator.SessionRunHook instead.

W0627 11:36:47.010504 139917697820480 deprecation_wrapper.py:119] From /private/home/yuxinwu/.local/lib/python3.6/site-packages/byteps/tensorflow/__init__.py:111: The name tf.train.Optimizer is deprecated. Please use tf.compat.v1.train.Optimizer instead.


In [3]: bps.push_pull(tf.constant([0.0]))
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-3-fc0e9eb4f2bb> in <module>()
----> 1 bps.push_pull(tf.constant([0.0]))

~/.local/lib/python3.6/site-packages/byteps/tensorflow/__init__.py in push_pull(tensor, scope, average, device_dense, device_sparse, compression)
50         byteps_size = tf.cast(size(), dtype=tensor.dtype)
51         tensor_compressed, ctx = compression.compress(tensor)
---> 52         summed_tensor_compressed = _push_pull(tensor_compressed, scope)
53         summed_tensor = compression.decompress(summed_tensor_compressed, ctx)
54         new_tensor = (tf.div(summed_tensor, byteps_size)

~/.local/lib/python3.6/site-packages/byteps/tensorflow/ops.py in _push_pull(tensor, scope, name)
80     if name is None and not _executing_eagerly():
81         name = 'BytePSPushPull_%s' % _normalize_name(tensor.name)
---> 82     TF_LIB_CTYPES.byteps_tensorflow_declare_tensor(ctypes.c_char_p(scope+name))
83     return C_LIB.byteps_push_pull(tensor, name=name)
84

TypeError: bytes or integer address expected instead of str instance

# Fixed code
def main():
    opt = parse_args()

    filehandler = logging.FileHandler(opt.logging_file)
    streamhandler = logging.StreamHandler()

    logger = logging.getLogger("")
    logger.setLevel(logging.INFO)
    logger.addHandler(filehandler)
    logger.addHandler(streamhandler)

    logger.info(opt)

    bps.init()

    batch_size = opt.batch_size
    classes = 1000
    num_training_samples = 1281167

    num_gpus = opt.num_gpus
    # batch_size *= max(1, num_gpus)
    context = mx.gpu(bps.local_rank()) if num_gpus > 0 else mx.cpu(bps.local_rank())
    num_workers = opt.num_workers
    nworker = bps.size()
    rank = bps.rank()

    lr_decay = opt.lr_decay
    lr_decay_period = opt.lr_decay_period
    if opt.lr_decay_period > 0:
        lr_decay_epoch = list(range(lr_decay_period, opt.num_epochs, lr_decay_period))
    else:
        lr_decay_epoch = [int(i) for i in opt.lr_decay_epoch.split(",")]
    lr_decay_epoch = [e - opt.warmup_epochs for e in lr_decay_epoch]
    num_batches = num_training_samples // (batch_size * nworker)

    lr_scheduler = LRSequential(
        [
            LRScheduler(
                "linear",
                base_lr=opt.warmup_lr,
                target_lr=opt.lr * nworker / bps.local_size(),
                nepochs=opt.warmup_epochs,
                iters_per_epoch=num_batches,
            ),
            LRScheduler(
                opt.lr_mode,
                base_lr=opt.lr * nworker / bps.local_size(),
                target_lr=0,
                nepochs=opt.num_epochs - opt.warmup_epochs,
                iters_per_epoch=num_batches,
                step_epoch=lr_decay_epoch,
                step_factor=lr_decay,
                power=2,
            ),
        ]
    )

    model_name = opt.model

    kwargs = {"ctx": context, "pretrained": opt.use_pretrained, "classes": classes}
    if opt.use_gn:
        from gluoncv.nn import GroupNorm

        kwargs["norm_layer"] = GroupNorm
    if model_name.startswith("vgg"):
        kwargs["batch_norm"] = opt.batch_norm
    elif model_name.startswith("resnext"):
        kwargs["use_se"] = opt.use_se

    if opt.last_gamma:
        kwargs["last_gamma"] = True

    if opt.compressor:
        optimizer = "sgd"
    else:
        optimizer = "nag"

    optimizer_params = {
        "wd": opt.wd,
        "momentum": opt.momentum,
        "lr_scheduler": lr_scheduler,
    }

    if opt.dtype != "float32":
        optimizer_params["multi_precision"] = True

    net = get_model(model_name, **kwargs)
    net.cast(opt.dtype)
    if opt.resume_params is not "":
        net.load_parameters(opt.resume_params, ctx=context)

    # teacher model for distillation training
    if opt.teacher is not None and opt.hard_weight < 1.0:
        teacher_name = opt.teacher
        teacher = get_model(teacher_name, pretrained=True, classes=classes, ctx=context)
        teacher.cast(opt.dtype)
        distillation = True
    else:
        distillation = False

    # Two functions for reading data from record file or raw images
    def get_data_rec(
        rec_train, rec_train_idx, rec_val, rec_val_idx, batch_size, num_workers
    ):
        rec_train = os.path.expanduser(rec_train)
        rec_train_idx = os.path.expanduser(rec_train_idx)
        rec_val = os.path.expanduser(rec_val)
        rec_val_idx = os.path.expanduser(rec_val_idx)
        jitter_param = 0.4
        lighting_param = 0.1
        input_size = opt.input_size
        crop_ratio = opt.crop_ratio if opt.crop_ratio > 0 else 0.875
        resize = int(math.ceil(input_size / crop_ratio))
        mean_rgb = [123.68, 116.779, 103.939]
        std_rgb = [58.393, 57.12, 57.375]

        def batch_fn(batch, ctx):
            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
            )
            return data, label

        train_data = mx.io.ImageRecordIter(
            path_imgrec=rec_train,
            path_imgidx=rec_train_idx,
            preprocess_threads=num_workers,
            shuffle=True,
            batch_size=batch_size,
            data_shape=(3, input_size, input_size),
            mean_r=mean_rgb[0],
            mean_g=mean_rgb[1],
            mean_b=mean_rgb[2],
            std_r=std_rgb[0],
            std_g=std_rgb[1],
            std_b=std_rgb[2],
            rand_mirror=True,
            random_resized_crop=True,
            max_aspect_ratio=4.0 / 3.0,
            min_aspect_ratio=3.0 / 4.0,
            max_random_area=1,
            min_random_area=0.08,
            brightness=jitter_param,
            saturation=jitter_param,
            contrast=jitter_param,
            pca_noise=lighting_param,
            num_parts=nworker,
            part_index=rank,
        )
        val_data = mx.io.ImageRecordIter(
            path_imgrec=rec_val,
            path_imgidx=rec_val_idx,
            preprocess_threads=num_workers,
            shuffle=False,
            batch_size=batch_size,
            resize=resize,
            data_shape=(3, input_size, input_size),
            mean_r=mean_rgb[0],
            mean_g=mean_rgb[1],
            mean_b=mean_rgb[2],
            std_r=std_rgb[0],
            std_g=std_rgb[1],
            std_b=std_rgb[2],
            num_parts=nworker,
            part_index=rank,
        )
        return train_data, val_data, batch_fn

    def get_data_loader(data_dir, batch_size, num_workers):
        normalize = transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        jitter_param = 0.4
        lighting_param = 0.1
        input_size = opt.input_size
        crop_ratio = opt.crop_ratio if opt.crop_ratio > 0 else 0.875
        resize = int(math.ceil(input_size / crop_ratio))

        def batch_fn(batch, ctx):
            data = gluon.utils.split_and_load(batch[0], ctx_list=ctx, batch_axis=0)
            label = gluon.utils.split_and_load(batch[1], ctx_list=ctx, batch_axis=0)
            return data, label

        transform_train = transforms.Compose(
            [
                transforms.RandomResizedCrop(input_size),
                transforms.RandomFlipLeftRight(),
                transforms.RandomColorJitter(
                    brightness=jitter_param,
                    contrast=jitter_param,
                    saturation=jitter_param,
                ),
                transforms.RandomLighting(lighting_param),
                transforms.ToTensor(),
                normalize,
            ]
        )
        transform_test = transforms.Compose(
            [
                transforms.Resize(resize, keep_ratio=True),
                transforms.CenterCrop(input_size),
                transforms.ToTensor(),
                normalize,
            ]
        )

        train_data = gluon.data.DataLoader(
            imagenet.classification.ImageNet(data_dir, train=True).transform_first(
                transform_train
            ),
            batch_size=batch_size,
            shuffle=True,
            last_batch="discard",
            num_workers=num_workers,
        )
        val_data = gluon.data.DataLoader(
            imagenet.classification.ImageNet(data_dir, train=False).transform_first(
                transform_test
            ),
            batch_size=batch_size,
            shuffle=False,
            num_workers=num_workers,
        )

        return train_data, val_data, batch_fn

    if opt.use_rec:
        train_data, val_data, batch_fn = get_data_rec(
            opt.rec_train,
            opt.rec_train_idx,
            opt.rec_val,
            opt.rec_val_idx,
            batch_size,
            num_workers,
        )
    else:
        train_data, val_data, batch_fn = get_data_loader(
            opt.data_dir, batch_size, num_workers
        )

    if opt.mixup:
        train_metric = mx.metric.RMSE()
    else:
        train_metric = mx.metric.Accuracy()
    acc_top1 = mx.metric.Accuracy()
    acc_top5 = mx.metric.TopKAccuracy(5)

    save_frequency = opt.save_frequency
    if opt.save_dir and save_frequency:
        save_dir = opt.save_dir
        makedirs(save_dir)
    else:
        save_dir = ""
        save_frequency = 0

    def mixup_transform(label, classes, lam=1, eta=0.0):
        if isinstance(label, nd.NDArray):
            label = [label]
        res = []
        for l in label:
            y1 = l.one_hot(
                classes, on_value=1 - eta + eta / classes, off_value=eta / classes
            )
            y2 = l[::-1].one_hot(
                classes, on_value=1 - eta + eta / classes, off_value=eta / classes
            )
            res.append(lam * y1 + (1 - lam) * y2)
        return res

    def smooth(label, classes, eta=0.1):
        if isinstance(label, nd.NDArray):
            label = [label]
        smoothed = []
        for l in label:
            res = l.one_hot(
                classes, on_value=1 - eta + eta / classes, off_value=eta / classes
            )
            smoothed.append(res)
        return smoothed

    def test(ctx, val_data):
        if opt.use_rec:
            val_data.reset()
        acc_top1.reset()
        acc_top5.reset()
        for i, batch in enumerate(val_data):
            data, label = batch_fn(batch, ctx)
            outputs = [net(X.astype(opt.dtype, copy=False)) for X in data]
            acc_top1.update(label, outputs)
            acc_top5.update(label, outputs)

        _, top1 = acc_top1.get()
        _, top5 = acc_top5.get()
        return (1 - top1, 1 - top5)

    def train(ctx):
        if isinstance(ctx, mx.Context):
            ctx = [ctx]
        if opt.resume_params is "":
            net.initialize(mx.init.MSRAPrelu(), ctx=ctx)

        if opt.no_wd:
            for k, v in net.collect_params(".*beta|.*gamma|.*bias").items():
                v.wd_mult = 0.0

        compression_params = {
            "compressor": opt.compressor,
            "ef": opt.ef,
            "momentum": opt.compress_momentum,
            "scaling": opt.onebit_scaling,
            "k": opt.k,
        }

        trainer = bps.DistributedTrainer(
            net.collect_params(),
            optimizer,
            optimizer_params,
            compression_params=compression_params,
        )

        if opt.resume_states is not "":
            trainer.load_states(opt.resume_states)

        if opt.label_smoothing or opt.mixup:
            sparse_label_loss = False
        else:
            sparse_label_loss = True
        if distillation:
            L = gcv.loss.DistillationSoftmaxCrossEntropyLoss(
                temperature=opt.temperature,
                hard_weight=opt.hard_weight,
                sparse_label=sparse_label_loss,
            )
        else:
            L = gluon.loss.SoftmaxCrossEntropyLoss(sparse_label=sparse_label_loss)

        best_val_score = 1

        for epoch in range(opt.resume_epoch, opt.num_epochs):
            tic = time.time()
            if opt.use_rec:
                train_data.reset()
            train_metric.reset()
            btic = time.time()

            for i, batch in enumerate(train_data):
                data, label = batch_fn(batch, ctx)

                if opt.mixup:
                    lam = np.random.beta(opt.mixup_alpha, opt.mixup_alpha)
                    if epoch >= opt.num_epochs - opt.mixup_off_epoch:
                        lam = 1
                    data = [lam * X + (1 - lam) * X[::-1] for X in data]

                    if opt.label_smoothing:
                        eta = 0.1
                    else:
                        eta = 0.0
                    label = mixup_transform(label, classes, lam, eta)

                elif opt.label_smoothing:
                    hard_label = label
                    label = smooth(label, classes)

                if distillation:
                    teacher_prob = [
                        nd.softmax(
                            teacher(X.astype(opt.dtype, copy=False)) / opt.temperature
                        )
                        for X in data
                    ]

                with ag.record():
                    outputs = [net(X.astype(opt.dtype, copy=False)) for X in data]
                    if distillation:
                        loss = [
                            L(
                                yhat.astype("float32", copy=False),
                                y.astype("float32", copy=False),
                                p.astype("float32", copy=False),
                            )
                            for yhat, y, p in zip(outputs, label, teacher_prob)
                        ]
                    else:
                        loss = [
                            L(yhat, y.astype(opt.dtype, copy=False))
                            for yhat, y in zip(outputs, label)
                        ]
                for l in loss:
                    l.backward()
                trainer.step(batch_size)

                if opt.mixup:
                    output_softmax = [
                        nd.SoftmaxActivation(out.astype("float32", copy=False))
                        for out in outputs
                    ]
                    train_metric.update(label, output_softmax)
                else:
                    if opt.label_smoothing:
                        train_metric.update(hard_label, outputs)
                    else:
                        train_metric.update(label, outputs)

                if opt.log_interval and not (i + 1) % opt.log_interval:
                    train_metric_name, train_metric_score = train_metric.get()
                    logger.info(
                        "Epoch[%d] Batch [%d]\tSpeed: %f samples/sec\t%s=%f\tlr=%f\ttime=%f"
                        % (
                            epoch,
                            i,
                            batch_size
                            * nworker
                            * opt.log_interval
                            / (time.time() - btic),
                            train_metric_name,
                            train_metric_score,
                            trainer.learning_rate,
                            time.time() - btic,
                        )
                    )
                    btic = time.time()

            train_metric_name, train_metric_score = train_metric.get()
            throughput = int(batch_size * nworker * i / (time.time() - tic))

            logger.info(
                "[Epoch %d] training: %s=%f"
                % (epoch, train_metric_name, train_metric_score)
            )
            logger.info(
                "[Epoch %d] speed: %d samples/sec\ttime cost: %f"
                % (epoch, throughput, time.time() - tic)
            )

            err_top1_val, err_top5_val = test(ctx, val_data)

            logger.info(
                "[Epoch %d] validation: err-top1=%f err-top5=%f"
                % (epoch, err_top1_val, err_top5_val)
            )

            if err_top1_val < best_val_score:
                best_val_score = err_top1_val
                net.save_parameters(
                    "%s/%.4f-imagenet-%s-%d-best.params"
                    % (save_dir, best_val_score, model_name, epoch)
                )
                trainer.save_states(
                    "%s/%.4f-imagenet-%s-%d-best.states"
                    % (save_dir, best_val_score, model_name, epoch)
                )

            if save_frequency and save_dir and (epoch + 1) % save_frequency == 0:
                net.save_parameters(
                    "%s/imagenet-%s-%d.params" % (save_dir, model_name, epoch)
                )
                trainer.save_states(
                    "%s/imagenet-%s-%d.states" % (save_dir, model_name, epoch)
                )

        if save_frequency and save_dir:
            net.save_parameters(
                "%s/imagenet-%s-%d.params" % (save_dir, model_name, opt.num_epochs - 1)
            )
            trainer.save_states(
                "%s/imagenet-%s-%d.states" % (save_dir, model_name, opt.num_epochs - 1)
            )

    if opt.mode == "hybrid":
        net.hybridize(static_alloc=True, static_shape=True)
        if distillation:
            teacher.hybridize(static_alloc=True, static_shape=True)
    train(context)