def main_worker(gpu, ngpus_per_node, args):
    args.gpu = gpu
    args.rank = ((args.rank * ngpus_per_node) + gpu)
    args.model = '{}-{}'.format(args.arch, os.path.splitext(os.path.basename(args.config_file))[0])
    if (args.gpu == 0):
        print('model:', args.model)
    print('rank: {} / {}'.format(args.rank, args.world_size))
    dist.init_process_group(backend=args.dist_backend, init_method=args.dist_url, world_size=args.world_size, rank=args.rank)
    torch.cuda.set_device(args.gpu)
    global best_pred, acclist_train, acclist_val
    if (args.gpu == 0):
        print(args)
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed(args.seed)
    cudnn.benchmark = True
    (transform_train, transform_val) = encoding.transforms.get_transform(args.dataset, args.base_size, args.crop_size)
    if (args.auto_policy is not None):
        print(f'Using auto_policy: {args.auto_policy}')
        from augment import Augmentation
        auto_policy = Augmentation(at.load(args.auto_policy))
        transform_train.transforms.insert(0, auto_policy)
    trainset = encoding.datasets.get_dataset(args.dataset, root=args.data_dir, transform=transform_train, train=True, download=True)
    valset = encoding.datasets.get_dataset(args.dataset, root=args.data_dir, transform=transform_val, train=False, download=True)
    train_sampler = torch.utils.data.distributed.DistributedSampler(trainset)
    train_loader = torch.utils.data.DataLoader(trainset, batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=True, sampler=train_sampler)
    val_sampler = torch.utils.data.distributed.DistributedSampler(valset, shuffle=False)
    val_loader = torch.utils.data.DataLoader(valset, batch_size=args.test_batch_size, shuffle=False, num_workers=args.workers, pin_memory=True, sampler=val_sampler)
    arch = importlib.import_module(('arch.' + args.arch))
    model = arch.config_network(args.config_file)
    if (args.gpu == 0):
        print(model)
    if (args.mixup > 0):
        train_loader = MixUpWrapper(args.mixup, 1000, train_loader, args.gpu)
        criterion = NLLMultiLabelSmooth(args.label_smoothing)
    elif (args.label_smoothing > 0.0):
        criterion = LabelSmoothing(args.label_smoothing)
    else:
        criterion = nn.CrossEntropyLoss()
    model.cuda(args.gpu)
    criterion.cuda(args.gpu)
    if args.no_bn_wd:
        parameters = model.named_parameters()
        param_dict = {}
        for (k, v) in parameters:
            param_dict[k] = v
        bn_params = [v for (n, v) in param_dict.items() if (('bn' in n) or ('bias' in n))]
        rest_params = [v for (n, v) in param_dict.items() if (not (('bn' in n) or ('bias' in n)))]
        if (args.gpu == 0):
            print(' Weight decay NOT applied to BN parameters ')
            print(f'len(parameters): {len(list(model.parameters()))} = {len(bn_params)} + {len(rest_params)}')
        optimizer = torch.optim.SGD([{'params': bn_params, 'weight_decay': 0}, {'params': rest_params, 'weight_decay': args.wd}], lr=args.lr, momentum=args.momentum, weight_decay=args.wd)
    else:
        optimizer = torch.optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.wd)
    if args.amp:
        (model, optimizer) = amp.initialize(model, optimizer, opt_level='O2')
        DDP = apex.parallel.DistributedDataParallel
        model = DDP(model, delay_allreduce=True)
    else:
        DDP = DistributedDataParallel
        model = DDP(model, device_ids=[args.gpu])
    if (args.resume is not None):
        if os.path.isfile(args.resume):
            if (args.gpu == 0):
                print("=> loading checkpoint '{}'".format(args.resume))
            checkpoint = torch.load(args.resume)
            args.start_epoch = ((checkpoint['epoch'] + 1) if (args.start_epoch == 0) else args.start_epoch)
            best_pred = checkpoint['best_pred']
            acclist_train = checkpoint['acclist_train']
            acclist_val = checkpoint['acclist_val']
            model.module.load_state_dict(checkpoint['state_dict'])
            optimizer.load_state_dict(checkpoint['optimizer'])
            if args.amp:
                amp.load_state_dict(checkpoint['amp'])
            if (args.gpu == 0):
                print("=> loaded checkpoint '{}' (epoch {})".format(args.resume, checkpoint['epoch']))
        else:
            raise RuntimeError("=> no resume checkpoint found at '{}'".format(args.resume))
    scheduler = LR_Scheduler(args.lr_scheduler, base_lr=args.lr, num_epochs=args.epochs, iters_per_epoch=len(train_loader), warmup_epochs=args.warmup_epochs)

    def train(epoch):
        train_sampler.set_epoch(epoch)
        model.train()
        losses = AverageMeter()
        top1 = AverageMeter()
        global best_pred, acclist_train
        tic = time.time()
        for (batch_idx, (data, target)) in enumerate(train_loader):
            scheduler(optimizer, batch_idx, epoch, best_pred)
            if (not args.mixup):
                (data, target) = (data.cuda(args.gpu), target.cuda(args.gpu))
            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, target)
            if args.amp:
                with amp.scale_loss(loss, optimizer) as scaled_loss:
                    scaled_loss.backward()
            else:
                loss.backward()
            optimizer.step()
            if (not args.mixup):
                acc1 = accuracy(output, target, topk=(1,))
                top1.update(acc1[0], data.size(0))
            losses.update(loss.item(), data.size(0))
            if (((batch_idx % 100) == 0) and (args.gpu == 0)):
                iter_per_sec = ((100.0 / (time.time() - tic)) if (batch_idx != 0) else (1.0 / (time.time() - tic)))
                tic = time.time()
                if args.mixup:
                    print('Epoch: {}, Iter: {}, Speed: {:.3f} iter/sec, Train loss: {:.3f}'.format(epoch, batch_idx, iter_per_sec, losses.avg.item()))
                else:
                    print('Epoch: {}, Iter: {}, Speed: {:.3f} iter/sec, Top1: {:.3f}'.format(epoch, batch_idx, iter_per_sec, top1.avg.item()))
        acclist_train += [top1.avg]

    def validate(epoch):
        model.eval()
        top1 = AverageMeter()
        top5 = AverageMeter()
        global best_pred, acclist_train, acclist_val
        is_best = False
        for (batch_idx, (data, target)) in enumerate(val_loader):
            (data, target) = (data.cuda(args.gpu), target.cuda(args.gpu))
            with torch.no_grad():
                output = model(data)
                (acc1, acc5) = accuracy(output, target, topk=(1, 5))
                top1.update(acc1[0], data.size(0))
                top5.update(acc5[0], data.size(0))
        (sum1, cnt1, sum5, cnt5) = torch_dist_sum(args.gpu, top1.sum, top1.count, top5.sum, top5.count)
        if args.eval:
            if (args.gpu == 0):
                top1_acc = (sum(sum1) / sum(cnt1))
                top5_acc = (sum(sum5) / sum(cnt5))
                print(('Validation: Top1: %.3f | Top5: %.3f' % (top1_acc, top5_acc)))
            return
        if (args.gpu == 0):
            top1_acc = (sum(sum1) / sum(cnt1))
            top5_acc = (sum(sum5) / sum(cnt5))
            print(('Validation: Top1: %.3f | Top5: %.3f' % (top1_acc, top5_acc)))
            acclist_val += [top1_acc]
            if (top1_acc > best_pred):
                best_pred = top1_acc
                is_best = True
            state_dict = {'epoch': epoch, 'state_dict': model.module.state_dict(), 'optimizer': optimizer.state_dict(), 'best_pred': best_pred, 'acclist_train': acclist_train, 'acclist_val': acclist_val}
            if args.amp:
                state_dict['amp'] = amp.state_dict()
            encoding.utils.save_checkpoint(state_dict, args=args, is_best=is_best)
    if args.export:
        if (args.gpu == 0):
            torch.save(model.module.state_dict(), (args.export + '.pth'))
        return
    if args.eval:
        validate(args.start_epoch)
        return
    for epoch in range(args.start_epoch, args.epochs):
        tic = time.time()
        train(epoch)
        if ((epoch % 10) == 0):
            validate(epoch)
        elapsed = (time.time() - tic)
        if (args.gpu == 0):
            print(f'Epoch: {epoch}, Time cost: {elapsed}')
    if (args.gpu == 0):
        encoding.utils.save_checkpoint({'epoch': (args.epochs - 1), 'state_dict': model.module.state_dict(), 'optimizer': optimizer.state_dict(), 'best_pred': best_pred, 'acclist_train': acclist_train, 'acclist_val': acclist_val}, args=args, is_best=False)
def main_worker(gpu, ngpus_per_node, args):
    args.gpu = gpu
    args.rank = ((args.rank * ngpus_per_node) + gpu)
    args.model = '{}-{}'.format(args.arch, os.path.splitext(os.path.basename(args.config_file))[0])
    if (args.gpu == 0):
        print('model:', args.model)
    print('rank: {} / {}'.format(args.rank, args.world_size))
    dist.init_process_group(backend=args.dist_backend, init_method=args.dist_url, world_size=args.world_size, rank=args.rank)
    torch.cuda.set_device(args.gpu)
    global best_pred, acclist_train, acclist_val
    if (args.gpu == 0):
        print(args)
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed(args.seed)
    cudnn.benchmark = True
    (transform_train, transform_val) = encoding.transforms.get_transform(args.dataset, args.base_size, args.crop_size)
    if (args.auto_policy is not None):
        print(f'Using auto_policy: {args.auto_policy}')
        from augment import Augmentation
        auto_policy = Augmentation(at.load(args.auto_policy))
        transform_train.transforms.insert(0, auto_policy)
    trainset = encoding.datasets.get_dataset(args.dataset, root=args.data_dir, transform=transform_train, train=True, download=True)
    valset = encoding.datasets.get_dataset(args.dataset, root=args.data_dir, transform=transform_val, train=False, download=True)
    train_sampler = torch.utils.data.distributed.DistributedSampler(trainset)
    train_loader = torch.utils.data.DataLoader(trainset, batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=True, sampler=train_sampler)
    val_sampler = torch.utils.data.distributed.DistributedSampler(valset, shuffle=False)
    val_loader = torch.utils.data.DataLoader(valset, batch_size=args.test_batch_size, shuffle=False, num_workers=args.workers, pin_memory=True, sampler=val_sampler)
    arch = importlib.import_module(('arch.' + args.arch))
    model = arch.config_network(args.config_file)
    if (args.gpu == 0):
        print(model)
    if (args.mixup > 0):
        train_loader = MixUpWrapper(args.mixup, 1000, train_loader, args.gpu)
        criterion = NLLMultiLabelSmooth(args.label_smoothing)
    elif (args.label_smoothing > 0.0):
        criterion = LabelSmoothing(args.label_smoothing)
    else:
        criterion = nn.CrossEntropyLoss()
    model.cuda(args.gpu)
    criterion.cuda(args.gpu)
    if args.no_bn_wd:
        parameters = model.named_parameters()
        param_dict = {}
        for (k, v) in parameters:
            param_dict[k] = v
        bn_params = [v for (n, v) in param_dict.items() if (('bn' in n) or ('bias' in n))]
        rest_params = [v for (n, v) in param_dict.items() if (not (('bn' in n) or ('bias' in n)))]
        if (args.gpu == 0):
            print(' Weight decay NOT applied to BN parameters ')
            print(f'len(parameters): {len(list(model.parameters()))} = {len(bn_params)} + {len(rest_params)}')
        optimizer = torch.optim.SGD([{'params': bn_params, 'weight_decay': 0}, {'params': rest_params, 'weight_decay': args.wd}], lr=args.lr, momentum=args.momentum, weight_decay=args.wd)
    else:
        optimizer = torch.optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.wd)
    if args.amp:
        (model, optimizer) = amp.initialize(model, optimizer, opt_level='O2')
        DDP = apex.parallel.DistributedDataParallel
        model = DDP(model, delay_allreduce=True)
    else:
        DDP = DistributedDataParallel
        model = DDP(model, device_ids=[args.gpu])
    if (args.resume is not None):
        if os.path.isfile(args.resume):
            if (args.gpu == 0):
                print("=> loading checkpoint '{}'".format(args.resume))
            checkpoint = torch.load(args.resume)
            args.start_epoch = ((checkpoint['epoch'] + 1) if (args.start_epoch == 0) else args.start_epoch)
            best_pred = checkpoint['best_pred']
            acclist_train = checkpoint['acclist_train']
            acclist_val = checkpoint['acclist_val']
            model.module.load_state_dict(checkpoint['state_dict'])
            optimizer.load_state_dict(checkpoint['optimizer'])
            if args.amp:
                amp.load_state_dict(checkpoint['amp'])
            if (args.gpu == 0):
                print("=> loaded checkpoint '{}' (epoch {})".format(args.resume, checkpoint['epoch']))
        else:
            raise RuntimeError("=> no resume checkpoint found at '{}'".format(args.resume))
    scheduler = LR_Scheduler(args.lr_scheduler, base_lr=args.lr, num_epochs=args.epochs, iters_per_epoch=len(train_loader), warmup_epochs=args.warmup_epochs)

    def train(epoch):
        train_sampler.set_epoch(epoch)
        model.train()
        losses = AverageMeter()
        top1 = AverageMeter()
        global best_pred, acclist_train
        tic = time.time()
        for (batch_idx, (data, target)) in enumerate(train_loader):
            scheduler(optimizer, batch_idx, epoch, best_pred)
            if (not args.mixup):
                (data, target) = (data.cuda(args.gpu), target.cuda(args.gpu))
            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, target)
            if args.amp:
                with amp.scale_loss(loss, optimizer) as scaled_loss:
                    scaled_loss.backward()
            else:
                loss.backward()
            optimizer.step()
            if (not args.mixup):
                acc1 = accuracy(output, target, topk=(1,))
                top1.update(acc1[0], data.size(0))
            losses.update(loss.item(), data.size(0))
            if (((batch_idx % 100) == 0) and (args.gpu == 0)):
                iter_per_sec = ((100.0 / (time.time() - tic)) if (batch_idx != 0) else (1.0 / (time.time() - tic)))
                tic = time.time()
                if args.mixup:
                    print('Epoch: {}, Iter: {}, Speed: {:.3f} iter/sec, Train loss: {:.3f}'.format(epoch, batch_idx, iter_per_sec, losses.avg.item()))
                else:
                    print('Epoch: {}, Iter: {}, Speed: {:.3f} iter/sec, Top1: {:.3f}'.format(epoch, batch_idx, iter_per_sec, top1.avg.item()))
        acclist_train += [top1.avg]

    def validate(epoch):
        model.eval()
        top1 = AverageMeter()
        top5 = AverageMeter()
        global best_pred, acclist_train, acclist_val
        is_best = False
        for (batch_idx, (data, target)) in enumerate(val_loader):
            (data, target) = (data.cuda(args.gpu), target.cuda(args.gpu))
            with torch.no_grad():
                output = model(data)
                (acc1, acc5) = accuracy(output, target, topk=(1, 5))
                top1.update(acc1[0], data.size(0))
                top5.update(acc5[0], data.size(0))
        (sum1, cnt1, sum5, cnt5) = torch_dist_sum(args.gpu, top1.sum, top1.count, top5.sum, top5.count)
        if args.eval:
            if (args.gpu == 0):
                top1_acc = (sum(sum1) / sum(cnt1))
                top5_acc = (sum(sum5) / sum(cnt5))
                print(('Validation: Top1: %.3f | Top5: %.3f' % (top1_acc, top5_acc)))
            return
        if (args.gpu == 0):
            top1_acc = (sum(sum1) / sum(cnt1))
            top5_acc = (sum(sum5) / sum(cnt5))
            print(('Validation: Top1: %.3f | Top5: %.3f' % (top1_acc, top5_acc)))
            acclist_val += [top1_acc]
            if (top1_acc > best_pred):
                best_pred = top1_acc
                is_best = True
            state_dict = {'epoch': epoch, 'state_dict': model.module.state_dict(), 'optimizer': optimizer.state_dict(), 'best_pred': best_pred, 'acclist_train': acclist_train, 'acclist_val': acclist_val}
            if args.amp:
                state_dict['amp'] = amp.state_dict()
            encoding.utils.save_checkpoint(state_dict, args=args, is_best=is_best)
    if args.export:
        if (args.gpu == 0):
            torch.save(model.module.state_dict(), (args.export + '.pth'))
        return
    if args.eval:
        validate(args.start_epoch)
        return
    for epoch in range(args.start_epoch, args.epochs):
        tic = time.time()
        train(epoch)
        if ((epoch % 10) == 0):
            validate(epoch)
        elapsed = (time.time() - tic)
        if (args.gpu == 0):
            print(f'Epoch: {epoch}, Time cost: {elapsed}')
    if (args.gpu == 0):
        encoding.utils.save_checkpoint({'epoch': (args.epochs - 1), 'state_dict': model.module.state_dict(), 'optimizer': optimizer.state_dict(), 'best_pred': best_pred, 'acclist_train': acclist_train, 'acclist_val': acclist_val}, args=args, is_best=False)