import argparse
import csv
import os

import torch
import tqdm
from torch import distributed
from torch.utils import data
from torchvision import datasets
from torchvision import transforms

from nets import nn
from utils import util

data_dir = os.path.join('..', 'Dataset', 'IMAGENET')


def batch(images, target, model, criterion=None):
    images = images.cuda()
    target = target.cuda()
    if criterion:
        with torch.cuda.amp.autocast():
            loss = criterion(model(images), target)
        return loss
    else:
        return util.accuracy(model(images), target, top_k=(1, 5))


def train(args):
    epochs = 350
    batch_size = 288
    util.set_seeds(args.rank)
    model = nn.EfficientNet().cuda()
    lr = batch_size * torch.cuda.device_count() * 0.256 / 4096
    optimizer = nn.RMSprop(util.add_weight_decay(model), lr, 0.9, 1e-3, momentum=0.9)
    ema = nn.EMA(model)

    if args.distributed:
        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank])
    else:
        model = torch.nn.DataParallel(model)
    criterion = nn.CrossEntropyLoss().cuda()
    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])

    dataset = datasets.ImageFolder(os.path.join(data_dir, 'train'),
                                   transforms.Compose([util.RandomResize(),
                                                       transforms.ColorJitter(0.4, 0.4, 0.4),
                                                       transforms.RandomHorizontalFlip(),
                                                       util.RandomAugment(),
                                                       transforms.ToTensor(), normalize]))
    if args.distributed:
        sampler = torch.utils.data.distributed.DistributedSampler(dataset)
    else:
        sampler = None

    loader = data.DataLoader(dataset, batch_size, sampler=sampler, num_workers=8, pin_memory=True)

    scheduler = nn.StepLR(optimizer)
    amp_scale = torch.cuda.amp.GradScaler()
    with open(f'weights/{scheduler.__str__()}.csv', 'w') as f:
        if args.local_rank == 0:
            writer = csv.DictWriter(f, fieldnames=['epoch', 'acc@1', 'acc@5'])
            writer.writeheader()
        best_acc1 = 0
        for epoch in range(0, epochs):
            if args.distributed:
                sampler.set_epoch(epoch)
            if args.local_rank == 0:
                print(('\n' + '%10s' * 2) % ('epoch', 'loss'))
                bar = tqdm.tqdm(loader, total=len(loader))
            else:
                bar = loader
            model.train()
            for images, target in bar:
                loss = batch(images, target, model, criterion)
                optimizer.zero_grad()
                amp_scale.scale(loss).backward()
                amp_scale.step(optimizer)
                amp_scale.update()

                ema.update(model)
                torch.cuda.synchronize()
                if args.local_rank == 0:
                    bar.set_description(('%10s' + '%10.4g') % ('%g/%g' % (epoch + 1, epochs), loss))

            scheduler.step(epoch + 1)
            if args.local_rank == 0:
                acc1, acc5 = test(ema.model.eval())
                writer.writerow({'acc@1': str(f'{acc1:.3f}'),
                                 'acc@5': str(f'{acc5:.3f}'),
                                 'epoch': str(epoch + 1).zfill(3)})
                util.save_checkpoint({'state_dict': ema.model.state_dict()}, acc1 > best_acc1)
                best_acc1 = max(acc1, best_acc1)
    if args.distributed:
        torch.distributed.destroy_process_group()
    torch.cuda.empty_cache()


def test(model=None):
    if model is None:
        model = nn.EfficientNet()
        model.load_state_dict(torch.load('weights/best.pt', 'cpu')['state_dict'])
        model = model.cuda()
        model.eval()

    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])

    dataset = datasets.ImageFolder(os.path.join(data_dir, 'val'),
                                   transforms.Compose([transforms.Resize(416),
                                                       transforms.CenterCrop(384),
                                                       transforms.ToTensor(), normalize]))

    loader = data.DataLoader(dataset, 48, num_workers=os.cpu_count(), pin_memory=True)
    top1 = util.AverageMeter()
    top5 = util.AverageMeter()
    with torch.no_grad():
        for images, target in tqdm.tqdm(loader, ('%10s' * 2) % ('acc@1', 'acc@5')):
            acc1, acc5 = batch(images, target, model)
            torch.cuda.synchronize()
            top1.update(acc1.item(), images.size(0))
            top5.update(acc5.item(), images.size(0))
        acc1, acc5 = top1.avg, top5.avg
        print('%10.3g' * 2 % (acc1, acc5))
    if model is None:
        torch.cuda.empty_cache()
    else:
        return acc1, acc5


def print_parameters():
    model = nn.EfficientNet().eval()
    _ = model(torch.zeros(1, 3, 224, 224))
    params = sum(p.numel() for p in model.parameters())
    print(f'Number of parameters: {int(params)}')


def benchmark():
    shape = (1, 3, 384, 384)
    util.torch2onnx(nn.EfficientNet().export().eval(), shape)
    util.onnx2caffe()
    util.print_benchmark(shape)


def main():
    # python -m torch.distributed.launch --nproc_per_node=3 main.py --train
    parser = argparse.ArgumentParser()
    parser.add_argument("--local_rank", default=0, type=int)
    parser.add_argument('--benchmark', action='store_true')
    parser.add_argument('--train', action='store_true')
    parser.add_argument('--test', action='store_true')

    args = parser.parse_args()
    args.distributed = False
    args.rank = 0
    if 'WORLD_SIZE' in os.environ:
        args.distributed = int(os.environ['WORLD_SIZE']) > 1
    if args.distributed:
        torch.cuda.set_device(args.local_rank)
        torch.distributed.init_process_group(backend='nccl', init_method='env://')
        args.rank = torch.distributed.get_rank()
    if args.local_rank == 0:
        if not os.path.exists('weights'):
            os.makedirs('weights')
    if args.local_rank == 0:
        print_parameters()
    if args.benchmark:
        benchmark()
    if args.train:
        train(args)
    if args.test:
        test()


if __name__ == '__main__':
    main()

[((15, 11, 15, 52), 'os.path.join', 'os.path.join', ({(15, 24, 15, 28): '""".."""', (15, 30, 15, 39): '"""Dataset"""', (15, 41, 15, 51): '"""IMAGENET"""'}, {}), "('..', 'Dataset', 'IMAGENET')", False, 'import os\n'), ((32, 4, 32, 29), 'utils.util.set_seeds', 'util.set_seeds', ({(32, 19, 32, 28): 'args.rank'}, {}), '(args.rank)', False, 'from utils import util\n'), ((36, 10, 36, 23), 'nets.nn.EMA', 'nn.EMA', ({(36, 17, 36, 22): 'model'}, {}), '(model)', False, 'from nets import nn\n'), ((43, 16, 43, 91), 'torchvision.transforms.Normalize', 'transforms.Normalize', (), '', False, 'from torchvision import transforms\n'), ((56, 13, 56, 98), 'torch.utils.data.DataLoader', 'data.DataLoader', (), '', False, 'from torch.utils import data\n'), ((58, 16, 58, 36), 'nets.nn.StepLR', 'nn.StepLR', ({(58, 26, 58, 35): 'optimizer'}, {}), '(optimizer)', False, 'from nets import nn\n'), ((59, 16, 59, 43), 'torch.cuda.amp.GradScaler', 'torch.cuda.amp.GradScaler', ({}, {}), '()', False, 'import torch\n'), ((96, 4, 96, 28), 'torch.cuda.empty_cache', 'torch.cuda.empty_cache', ({}, {}), '()', False, 'import torch\n'), ((106, 16, 106, 91), 'torchvision.transforms.Normalize', 'transforms.Normalize', (), '', False, 'from torchvision import transforms\n'), ((114, 11, 114, 30), 'utils.util.AverageMeter', 'util.AverageMeter', ({}, {}), '()', False, 'from utils import util\n'), ((115, 11, 115, 30), 'utils.util.AverageMeter', 'util.AverageMeter', ({}, {}), '()', False, 'from utils import util\n'), ((140, 4, 140, 21), 'utils.util.onnx2caffe', 'util.onnx2caffe', ({}, {}), '()', False, 'from utils import util\n'), ((141, 4, 141, 31), 'utils.util.print_benchmark', 'util.print_benchmark', ({(141, 25, 141, 30): 'shape'}, {}), '(shape)', False, 'from utils import util\n'), ((146, 13, 146, 38), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ({}, {}), '()', False, 'import argparse\n'), ((35, 27, 35, 55), 'utils.util.add_weight_decay', 'util.add_weight_decay', ({(35, 49, 35, 54): 'model'}, {}), '(model)', False, 'from utils import util\n'), ((39, 16, 39, 94), 'torch.nn.parallel.DistributedDataParallel', 'torch.nn.parallel.DistributedDataParallel', (), '', False, 'import torch\n'), ((41, 16, 41, 44), 'torch.nn.DataParallel', 'torch.nn.DataParallel', ({(41, 38, 41, 43): 'model'}, {}), '(model)', False, 'import torch\n'), ((45, 35, 45, 66), 'os.path.join', 'os.path.join', ({(45, 48, 45, 56): 'data_dir', (45, 58, 45, 65): '"""train"""'}, {}), "(data_dir, 'train')", False, 'import os\n'), ((52, 18, 52, 74), 'torch.utils.data.distributed.DistributedSampler', 'torch.utils.data.distributed.DistributedSampler', ({(52, 66, 52, 73): 'dataset'}, {}), '(dataset)', False, 'import torch\n'), ((95, 8, 95, 49), 'torch.distributed.destroy_process_group', 'torch.distributed.destroy_process_group', ({}, {}), '()', False, 'import torch\n'), ((101, 16, 101, 33), 'nets.nn.EfficientNet', 'nn.EfficientNet', ({}, {}), '()', False, 'from nets import nn\n'), ((108, 35, 108, 64), 'os.path.join', 'os.path.join', ({(108, 48, 108, 56): 'data_dir', (108, 58, 108, 63): '"""val"""'}, {}), "(data_dir, 'val')", False, 'import os\n'), ((116, 9, 116, 24), 'torch.no_grad', 'torch.no_grad', ({}, {}), '()', False, 'import torch\n'), ((117, 30, 117, 82), 'tqdm.tqdm', 'tqdm.tqdm', ({(117, 40, 117, 46): 'loader', (117, 48, 117, 81): "('%10s' * 2 % ('acc@1', 'acc@5'))"}, {}), "(loader, '%10s' * 2 % ('acc@1', 'acc@5'))", False, 'import tqdm\n'), ((125, 8, 125, 32), 'torch.cuda.empty_cache', 'torch.cuda.empty_cache', ({}, {}), '()', False, 'import torch\n'), ((132, 14, 132, 41), 'torch.zeros', 'torch.zeros', ({(132, 26, 132, 27): '1', (132, 29, 132, 30): '3', (132, 32, 132, 35): '224', (132, 37, 132, 40): '224'}, {}), '(1, 3, 224, 224)', False, 'import torch\n'), ((158, 8, 158, 46), 'torch.cuda.set_device', 'torch.cuda.set_device', ({(158, 30, 158, 45): 'args.local_rank'}, {}), '(args.local_rank)', False, 'import torch\n'), ((159, 8, 159, 82), 'torch.distributed.init_process_group', 'torch.distributed.init_process_group', (), '', False, 'import torch\n'), ((160, 20, 160, 48), 'torch.distributed.get_rank', 'torch.distributed.get_rank', ({}, {}), '()', False, 'import torch\n'), ((22, 13, 22, 38), 'torch.cuda.amp.autocast', 'torch.cuda.amp.autocast', ({}, {}), '()', False, 'import torch\n'), ((33, 12, 33, 29), 'nets.nn.EfficientNet', 'nn.EfficientNet', ({}, {}), '()', False, 'from nets import nn\n'), ((42, 16, 42, 37), 'nets.nn.CrossEntropyLoss', 'nn.CrossEntropyLoss', ({}, {}), '()', False, 'from nets import nn\n'), ((62, 21, 62, 78), 'csv.DictWriter', 'csv.DictWriter', (), '', False, 'import csv\n'), ((113, 54, 113, 68), 'os.cpu_count', 'os.cpu_count', ({}, {}), '()', False, 'import os\n'), ((119, 12, 119, 36), 'torch.cuda.synchronize', 'torch.cuda.synchronize', ({}, {}), '()', False, 'import torch\n'), ((131, 12, 131, 29), 'nets.nn.EfficientNet', 'nn.EfficientNet', ({}, {}), '()', False, 'from nets import nn\n'), ((162, 15, 162, 40), 'os.path.exists', 'os.path.exists', ({(162, 30, 162, 39): '"""weights"""'}, {}), "('weights')", False, 'import os\n'), ((163, 12, 163, 34), 'os.makedirs', 'os.makedirs', ({(163, 24, 163, 33): '"""weights"""'}, {}), "('weights')", False, 'import os\n'), ((34, 22, 34, 47), 'torch.cuda.device_count', 'torch.cuda.device_count', ({}, {}), '()', False, 'import torch\n'), ((46, 55, 46, 74), 'utils.util.RandomResize', 'util.RandomResize', ({}, {}), '()', False, 'from utils import util\n'), ((47, 55, 47, 92), 'torchvision.transforms.ColorJitter', 'transforms.ColorJitter', ({(47, 78, 47, 81): '0.4', (47, 83, 47, 86): '0.4', (47, 88, 47, 91): '0.4'}, {}), '(0.4, 0.4, 0.4)', False, 'from torchvision import transforms\n'), ((48, 55, 48, 88), 'torchvision.transforms.RandomHorizontalFlip', 'transforms.RandomHorizontalFlip', ({}, {}), '()', False, 'from torchvision import transforms\n'), ((49, 55, 49, 75), 'utils.util.RandomAugment', 'util.RandomAugment', ({}, {}), '()', False, 'from utils import util\n'), ((50, 55, 50, 76), 'torchvision.transforms.ToTensor', 'transforms.ToTensor', ({}, {}), '()', False, 'from torchvision import transforms\n'), ((82, 16, 82, 40), 'torch.cuda.synchronize', 'torch.cuda.synchronize', ({}, {}), '()', False, 'import torch\n'), ((102, 30, 102, 66), 'torch.load', 'torch.load', ({(102, 41, 102, 58): '"""weights/best.pt"""', (102, 60, 102, 65): '"""cpu"""'}, {}), "('weights/best.pt', 'cpu')", False, 'import torch\n'), ((109, 55, 109, 77), 'torchvision.transforms.Resize', 'transforms.Resize', ({(109, 73, 109, 76): '416'}, {}), '(416)', False, 'from torchvision import transforms\n'), ((110, 55, 110, 81), 'torchvision.transforms.CenterCrop', 'transforms.CenterCrop', ({(110, 77, 110, 80): '384'}, {}), '(384)', False, 'from torchvision import transforms\n'), ((111, 55, 111, 76), 'torchvision.transforms.ToTensor', 'transforms.ToTensor', ({}, {}), '()', False, 'from torchvision import transforms\n'), ((139, 20, 139, 37), 'nets.nn.EfficientNet', 'nn.EfficientNet', ({}, {}), '()', False, 'from nets import nn\n')]