'''
References:
    https://github.com/TropComplique/trained-ternary-quantization
    MIT License - Copyright (c) 2017 <NAME>
    https://github.com/uoguelph-mlrg/Cutout
    Educational Community License, Version 2.0 (ECL-2.0) - Copyright (c) 2019 <NAME>
    https://github.com/lukemelas/EfficientNet-PyTorch
    Apache License, Version 2.0 - Copyright (c) 2019 <NAME>
    https://github.com/akamaster/pytorch_resnet_cifar10
    <NAME>'s ResNet implementation for CIFAR10/CIFAR100 in PyTorch

This document is the main document for applying grid search on depth and width scaling factors for the MicroNet
architectures which can be found in ./model/model.

This document contains the following functions/classes:

    - main()
        - Cutout(object)
        - grid_search(train_loader, val_loader, criterion, alpha, beta)
            - train(train_loader, model, criterion, optimizer, epoch)
            - evaluate(model, loss, val_iterator)

'''

import argparse
import time
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim
import torch.utils.data
import torchvision.transforms as transforms
import torchvision.datasets as datasets
#ssfrom torchsummary import summary
import numpy as np
import os
import PIL
from model import micronet, image_micronet, EfficientNet, best_cifar_micronet, lenet5


parser = argparse.ArgumentParser(description='Grid search for model compound scaling')
parser.add_argument('--workers', default=4, type=int, metavar='4',
                    help='Number of data loading workers (default: 4 per GPU)')
parser.add_argument('--epochs', default=200, type=int, metavar=200,
                    help='Number of epochs to run (default: 200)')
parser.add_argument('--start-epoch', default=0, type=int, metavar=0,
                    help='Start epoch, e.g. for restarts (default: 0)')
parser.add_argument('--batch-size', type=int, default=256, metavar=256,
                    help='Batch size for training (default=256)')
parser.add_argument('--val-batch-size', type=int, default=512, metavar=512,
                    help='Batch size for validation (default=512)')
parser.add_argument('--lr', default=0.1, type=float,
                    metavar=0.1, help='Initial learning rate (default 0.1)')
parser.add_argument('--momentum', default=0.9, type=float, metavar=0.9,
                    help='Momentum for SGD (default 0.9)')
parser.add_argument('--weight-decay', default=5e-4, type=float,
                    metavar=5e-4, help='Weight decay (default: 5e-4)')
parser.add_argument('--grid', type=float, default=(1.4, 1.2), nargs='+', metavar='1.4 1.2 1.2 1.3 ...',
                    help='Grid of alternating depth and width multipliers d and w (e.g.: d0, w0, d1, w1, ..., '
                         'dn, wn for n tuples)')
parser.add_argument('--phi', type=float, default=3.5, metavar=3.5,
                    help='Phi is the exponential scaling factor for width and depth multipliers (default: 3.5)')
parser.add_argument('--cuda', default=True, action='store_true',
                    help='By default CUDA training on GPU is enabled')
parser.add_argument('--no-cuda', dest='cuda', action='store_false',
                    help='Disables CUDA training and maps data plus model to CPU')
parser.add_argument('--dataset', default='CIFAR100', type=str, metavar='cifar',
                    help='Dataset to use. Choose from [CIFAR100, CIFAR10, MNIST, ImageNet] (default: CIFAR100)')
parser.add_argument('--image-size', default=32, type=int, metavar=32,
                    help='Input image size. Choose from [32 for CIFAR, 128-600 for ImageNet] (default: 32)')
parser.add_argument('--data-path', default='../data', type=str, metavar='/path',
                    help='Path to ImageNet data. CIFAR data will be downloaded to "../data" directory automatically '
                         'if the data-path argument is ignored')
parser.add_argument('--slurm-save', type=str, metavar='/path',
                    help='Path where slurm saves job results.')
parser.add_argument('--resume', default='./model_compound_scaling/saved_models/checkpoint.pt', type=str,
                    metavar='/path_cp',
                    help='Path to latest checkpoint. If it exists the train procedure will be resumed'
                         '(default: ./model_compound_scaling/saved_models/checkpoint.pt)')
parser.add_argument('--no-resume', dest='resume', action='store_false',
                    help='Do not resume from checkpoint')


def train(train_loader, model, criterion, optimizer, epoch):
    """
    Training procedure.

    Parameters:
    -----------
        train_loader:
            PyTorch Dataloader for given train dataset.
        model:
            The neural network model.
        criterion:
            Loss function to use (e.g. cross entropy, MSE, ...).
        optimizer:
            Updating model parameters with Gradient Descent plus Nesterov momentum and weight decay.
        epoch:
            Current training epoch.

    Returns:
    --------
        running loss
        running accuracy
    """

    # Resetting variables for next epoch
    running_loss = 0.0
    running_accuracy = 0.0
    n_steps = 0

    # Iterating over training dataloader
    for input, target in train_loader:

        # Measure data loading time
        end = time.time()

        # Resetting variables for calculating current batch accuracy
        correct = 0
        total = 0

        # Map data to GPU if available
        if use_cuda:
            input, target = input.cuda(), target.cuda(non_blocking=True)

        # Forward pass
        output = model(input)

        # Benchmark mode can allocate large memory blocks during the very first forward to test algorithms
        if epoch == 0:
            torch.cuda.empty_cache()

        # Compute batch_loss
        batch_loss = criterion(output, target)

        # Compute accuracy
        _, predicted = output.max(1)
        total += target.size(0)
        correct += predicted.eq(target).sum().item()
        batch_accuracy = 100. * correct / total

        # Compute gradient and do optimization step
        optimizer.zero_grad()
        batch_loss.backward()
        optimizer.step()

        # Summing up batch losses and accuracies over each step
        running_loss += batch_loss.float()
        running_accuracy += batch_accuracy
        n_steps += 1

        # Printing preliminary results
        if n_steps % 50 == 0:
            total_steps = int(len(train_loader.dataset) / train_loader.batch_size) + 1
            print('Epoch: {}, step: {}/{}, running_loss: {:.4f}, batch_acc: {:.4f}, running_acc: {:.4f}, '
                  'elapsed time: {:.2f} s, max_mem_alloc: {:.2f} GB, max_mem_cache {:.2f} GB'.format(
                epoch, n_steps, total_steps, running_loss / n_steps, batch_accuracy, running_accuracy / n_steps,
                                             (time.time() - end) * 100, torch.cuda.max_memory_allocated() / 1024 ** 3,
                                             torch.cuda.max_memory_cached() / 1024 ** 3))

        del output
        del batch_loss

    return running_loss / n_steps, running_accuracy / n_steps


def evaluate(model, loss, val_iterator):
    """
    Evaluating the model (no backward path / opimization step)

    Parameters:
    -----------
        model:
            The neural network model.
        loss:
            Loss function used for optimization (e.g. cross entropy, MSE, ...).
        val_iterator:
            PyTorch Dataloader for dataset which should be evaluated.

    Returns:
    --------
        Validation loss.
        Validation accuracy.
    """

    # Initializing parameters
    loss_value = 0.0
    accuracy = 0.0
    total_samples = 0

    with torch.no_grad():

        # Iterating over validation dataloader
        for data, labels in val_iterator:

            # Resetting variables for calculating current batch accuracy
            correct = 0
            total = 0

            # Map data to GPU if available
            if use_cuda:
                data = data.cuda()
                labels = labels.cuda(non_blocking=True)

            n_batch_samples = labels.size()[0]
            logits = model(data)

            # Compute batch loss
            batch_loss = loss(logits, labels)

            # Compute batch accuracy
            _, predicted = logits.max(1)
            total += labels.size(0)
            correct += predicted.eq(labels).sum().item()
            batch_accuracy = 100. * correct / total

            # Summing up batch losses and accuracies over each step
            loss_value += batch_loss.float() * n_batch_samples
            accuracy += batch_accuracy * n_batch_samples
            total_samples += n_batch_samples

            del logits
            del batch_loss

        return loss_value / total_samples, accuracy / total_samples


def grid_search(train_loader, val_loader, criterion, alpha, beta):
    """
    Builds the model with given scaling factors, sets up optimizer and learning rate schedulers plus executes
    training and evaluation of the model. A checkpoint is created each epoch. Also the best model will be saved.

    Parameters:
    -----------
        train_loader:
            PyTorch Dataloader for given train dataset.
        val_loader:
            PyTorch Dataloader for given validation dataset.
        criterion:
            Loss function to use (e.g. cross entropy, MSE, ...).
        alpha:
            Scaling factor for model depth.
        beta:
            Scaling factor for model width.
    """

    # Initializing training variables
    best_acc = 0
    all_losses = []

    # Initializing log file
    if args.slurm_save == None:
        logfile = open('./model_compound_scaling/logfiles/logfile.txt', 'a+')
    else:
        logfile = open(args.slurm_save + '/logfile.txt', 'a+')
    logfile.write('depth multiplier: {}, width multiplier: {}\n'.format(alpha, beta))

    # Building the model
    if args.dataset == 'CIFAR100':
        model = micronet(d_multiplier=alpha, w_multiplier=beta, num_classes=100)

    elif args.dataset == 'CIFAR10':
        model = micronet(d_multiplier=alpha, w_multiplier=beta, num_classes=10)

    elif args.dataset == 'ImageNet':
        model = image_micronet(d_multiplier=alpha, w_multiplier=beta)

    elif args.dataset == 'MNIST':
        model = lenet5(d_multiplier=alpha, w_multiplier=beta)

    # If multipile GPUs are used
    if use_cuda and torch.cuda.device_count() > 1:
        print("Let's use", torch.cuda.device_count(), "GPUs!")
        model = nn.DataParallel(model)

    # Transfers model to device (GPU/CPU). Device is globally initialized.
    model.to(device)

    # Defining the optimizer
    optimizer = torch.optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum,
                                weight_decay=args.weight_decay, nesterov=True)

    # KERAS like summary of the model architecture
    # summary(your_model, input_size=(channels, H, W), batch_size=-1, device="cuda")
    if use_cuda:
        if args.dataset == 'CIFAR100' or args.dataset == 'CIFAR10':
            #summary(model, (3, 32, 32), batch_size=args.batch_size)
            print(model)

        elif args.dataset == 'ImageNet':
            #summary(model, (3, args.image_size, args.image_size), batch_size=args.batch_size)
            print(model)

        elif args.dataset == 'MNIST':
            #summary(model, (3, args.image_size, args.image_size), batch_size=args.batch_size)
            print(model)

    # Optionally resume from a checkpoint
    if args.resume:
        if os.path.isfile(args.resume):
            print("=> loading checkpoint '{}'".format(args.resume))
            checkpoint = torch.load(args.resume)
            args.start_epoch = checkpoint['epoch']
            best_acc = checkpoint['acc']
            model.load_state_dict(checkpoint['state_dict'])
            optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
            load_last_epoch = checkpoint['epoch']-1
            print("=> loaded checkpoint '{}' (epoch {})".format(args.resume, checkpoint['epoch']))
        else:
            print("=> no checkpoint found at '{}'".format(args.resume))
            load_last_epoch = -1
    else:
        load_last_epoch = -1

    # Learning rate schedulers for cifar_micronet and imagenet_micronet
    lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer,
                                                           T_max = args.epochs,
                                                           eta_min = 0,
                                                           last_epoch = load_last_epoch)

    # START TRAINING
    start_time = time.time()
    model.train()

    for epoch in range(args.start_epoch, args.epochs):

        print('current lr {:.5e}'.format(optimizer.param_groups[0]['lr']))

        # Executing training process
        running_loss, running_accuracy = train(train_loader, model, criterion, optimizer, epoch)

        # Evaluation
        model.eval()
        val_loss, val_accuracy = evaluate(model, criterion, val_loader)

        # Logging the accuracies
        all_losses += [(epoch, running_loss, val_loss, running_accuracy, val_accuracy)]
        print('Epoch {0} running loss {1:.3f} val loss {2:.3f}  running acc {3:.3f} '
              'val acc{4:.3f}  time {5:.3f}'.format(*all_losses[-1], time.time() - start_time))
        logfile.write('Epoch {0} running loss {1:.3f} val loss {2:.3f}  running acc {3:.3f} '
              'val acc{4:.3f}  time {5:.3f}\n'.format(*all_losses[-1], time.time() - start_time))

        # Saving checkpoint
        if args.slurm_save == None:
            cp_savepath = './model_compound_scaling/saved_models/checkpoint.pt'
            best_savepath = './model_compound_scaling/saved_models/best_model.pt'
        else:
            cp_savepath = args.slurm_save + '/checkpoint.pt'
            best_savepath = args.slurm_save + '/best_model.pt'

        torch.save({
            'epoch': epoch,
            'state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'acc': val_accuracy,
            'lr': optimizer.param_groups[0]['lr']
        }, cp_savepath)

        # Make a lr scheduler step
        lr_scheduler.step()

        # Checking if current epoch yielded best validation accuracy
        is_best = val_accuracy > best_acc
        best_acc = max(val_accuracy, best_acc)

        # If so, saving best model state_dict
        if is_best and epoch > 0:
            torch.save(model.state_dict(), best_savepath)

        # Switch back to train mode
        model.train()
        start_time = time.time()


class Cutout(object):
    """
    Randomly mask out one or more patches from an image.

    Parameters:
    -----------
        n_holes (int):
            Number of patches to cut out of each image.
        length (int):
            The length (in pixels) of each square patch.
    """

    def __init__(self, n_holes, length):
        self.n_holes = n_holes
        self.length = length

    def __call__(self, img):
        """
        Parameters:
        -----------
            img (Tensor):
                Tensor image of size (C, H, W).
        Returns:
        --------
            Tensor:
                Image with n_holes of dimension length x length cut out of it.
        """
        h = img.size(1)
        w = img.size(2)

        mask = np.ones((h, w), np.float32)

        for n in range(self.n_holes):
            y = np.random.randint(h)
            x = np.random.randint(w)

            y1 = np.clip(y - self.length // 2, 0, h)
            y2 = np.clip(y + self.length // 2, 0, h)
            x1 = np.clip(x - self.length // 2, 0, w)
            x2 = np.clip(x + self.length // 2, 0, w)

            mask[y1: y2, x1: x2] = 0.

        mask = torch.from_numpy(mask)
        mask = mask.expand_as(img)
        img = img * mask

        return img


def main():
    """
    --------------------------------------------- MAIN --------------------------------------------------------
    Loads the data and executes the grid search on depth and width scaling factors.

    """
    # Manual seed for reproducibility
    torch.manual_seed(363636)

    # Global instances
    global args, use_cuda, device
    # Instantiating the parser
    args = parser.parse_args()
    # Global CUDA flag
    use_cuda = args.cuda and torch.cuda.is_available()
    # Defining device and device's map locationo
    device = torch.device("cuda" if use_cuda else "cpu")
    print('chosen device: ', device)

    # Defining loss function and printing CUDA information (if available)
    if use_cuda:
        print("PyTorch version: ")
        print(torch.__version__)
        print("CUDA Version: ")
        print(torch.version.cuda)
        print("cuDNN version is: ")
        print(cudnn.version())
        cudnn.benchmark = True
        criterion = nn.CrossEntropyLoss().cuda()
    else:
        criterion = nn.CrossEntropyLoss()

    # Dataloaders for CIFAR, ImageNet and MNIST
    if args.dataset == 'CIFAR100':

        normalize = transforms.Normalize(mean=[x / 255.0 for x in [125.3, 123.0, 113.9]],
                                         std=[x / 255.0 for x in [63.0, 62.1, 66.7]])

        kwargs = {'num_workers': args.workers, 'pin_memory': True} if use_cuda else {}

        train_loader = torch.utils.data.DataLoader(
            datasets.CIFAR100(root=args.data_path, train=True, transform=transforms.Compose([
                transforms.RandomHorizontalFlip(),
                transforms.RandomCrop(32, 4),
                transforms.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.3, hue=0.075),
                transforms.ToTensor(),
                normalize,
                Cutout(n_holes=1, length=16),
            ]), download=True),
            batch_size=args.batch_size, shuffle=True, **kwargs)

        val_loader = torch.utils.data.DataLoader(
            datasets.CIFAR100(root=args.data_path, train=False, transform=transforms.Compose([
                transforms.ToTensor(),
                normalize,
            ])),
            batch_size=args.val_batch_size, shuffle=False, **kwargs)

    elif args.dataset == 'ImageNet':

        traindir = os.path.join(args.data_path, 'train')
        valdir = os.path.join(args.data_path, 'val')
        normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                         std=[0.229, 0.224, 0.225])

        train_dataset = datasets.ImageFolder(
            traindir,
            transforms.Compose([
                transforms.RandomResizedCrop(args.image_size),
                transforms.RandomHorizontalFlip(),
                #transforms.RandomAffine(degrees=15),
                #transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.05),
                transforms.ToTensor(),
                normalize,
            ]))
        train_loader = torch.utils.data.DataLoader(
            train_dataset, batch_size=args.batch_size, shuffle=True,
            num_workers=args.workers, pin_memory=True)

        image_size = args.image_size
        val_dataset = datasets.ImageFolder(
            valdir,
            transforms.Compose([
                transforms.Resize(image_size, interpolation=PIL.Image.BICUBIC),
                transforms.CenterCrop(image_size),
                transforms.ToTensor(),
                normalize,
            ]))
        val_loader = torch.utils.data.DataLoader(
            val_dataset, batch_size=args.val_batch_size, shuffle=False,
            num_workers=args.workers, pin_memory=True)

    elif args.dataset == 'MNIST':

        kwargs = {'num_workers': args.workers, 'pin_memory': True} if use_cuda else {}

        train_loader = torch.utils.data.DataLoader(
            datasets.MNIST(args.data_path, train=True, download=True,
                           transform=transforms.Compose([
                               transforms.ToTensor(),
                               transforms.Normalize((0.1307,), (0.3081,))
                           ])),
            batch_size=args.batch_size, shuffle=True, **kwargs)
        val_loader = torch.utils.data.DataLoader(
            datasets.MNIST(args.data_path, train=False, transform=transforms.Compose([
                transforms.ToTensor(),
                transforms.Normalize((0.1307,), (0.3081,))
            ])),
            batch_size=args.val_batch_size, shuffle=True, **kwargs)

    elif args.dataset == 'CIFAR10':

        normalize = transforms.Normalize(mean=[x / 255.0 for x in [125.3, 123.0, 113.9]],
                                         std=[x / 255.0 for x in [63.0, 62.1, 66.7]])

        kwargs = {'num_workers': args.workers, 'pin_memory': True} if use_cuda else {}

        train_loader = torch.utils.data.DataLoader(
            datasets.CIFAR10(root=args.data_path, train=True, transform=transforms.Compose([
                transforms.RandomHorizontalFlip(),
                transforms.RandomCrop(32, 4),
                transforms.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.3, hue=0.075),
                transforms.ToTensor(),
                normalize,
                Cutout(n_holes=1, length=16),
            ]), download=True),
            batch_size=args.batch_size, shuffle=True, **kwargs)

        val_loader = torch.utils.data.DataLoader(
            datasets.CIFAR10(root=args.data_path, train=False, transform=transforms.Compose([
                transforms.ToTensor(),
                normalize,
            ])),
            batch_size=args.val_batch_size, shuffle=False, **kwargs)

    # original grid = [(1.0, 1.0), (1.9, 1.0), (1.7, 1.1), (1.6, 1.1), (1.4, 1.2), (1.2, 1.3), (1.0, 1.4)]

    grid = [(args.grid[i], args.grid[i+1]) for i in range(0, len(args.grid), 2)]

    for coeff in grid:
        alpha = coeff[0] ** args.phi
        beta = coeff[1] ** args.phi
        grid_search(train_loader, val_loader, criterion, alpha, beta)


if __name__ == '__main__':
    t1_start = time.perf_counter()
    main()
    t1_stop = time.perf_counter()
    print("Elapsed time: %.2f [s]" % (t1_stop - t1_start))

'''
References:
    https://github.com/TropComplique/trained-ternary-quantization
    MIT License - Copyright (c) 2017 <NAME>
    https://github.com/uoguelph-mlrg/Cutout
    Educational Community License, Version 2.0 (ECL-2.0) - Copyright (c) 2019 <NAME>
    https://github.com/lukemelas/EfficientNet-PyTorch
    Apache License, Version 2.0 - Copyright (c) 2019 <NAME>
    https://github.com/akamaster/pytorch_resnet_cifar10
    <NAME>'s ResNet implementation for CIFAR10/CIFAR100 in PyTorch

This document is the main document for applying grid search on depth and width scaling factors for the MicroNet
architectures which can be found in ./model/model.

This document contains the following functions/classes:

    - main()
        - Cutout(object)
        - grid_search(train_loader, val_loader, criterion, alpha, beta)
            - train(train_loader, model, criterion, optimizer, epoch)
            - evaluate(model, loss, val_iterator)

'''

import argparse
import time
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim
import torch.utils.data
import torchvision.transforms as transforms
import torchvision.datasets as datasets
#ssfrom torchsummary import summary
import numpy as np
import os
import PIL
from model import micronet, image_micronet, EfficientNet, best_cifar_micronet, lenet5


parser = argparse.ArgumentParser(description='Grid search for model compound scaling')
parser.add_argument('--workers', default=4, type=int, metavar='4',
                    help='Number of data loading workers (default: 4 per GPU)')
parser.add_argument('--epochs', default=200, type=int, metavar=200,
                    help='Number of epochs to run (default: 200)')
parser.add_argument('--start-epoch', default=0, type=int, metavar=0,
                    help='Start epoch, e.g. for restarts (default: 0)')
parser.add_argument('--batch-size', type=int, default=256, metavar=256,
                    help='Batch size for training (default=256)')
parser.add_argument('--val-batch-size', type=int, default=512, metavar=512,
                    help='Batch size for validation (default=512)')
parser.add_argument('--lr', default=0.1, type=float,
                    metavar=0.1, help='Initial learning rate (default 0.1)')
parser.add_argument('--momentum', default=0.9, type=float, metavar=0.9,
                    help='Momentum for SGD (default 0.9)')
parser.add_argument('--weight-decay', default=5e-4, type=float,
                    metavar=5e-4, help='Weight decay (default: 5e-4)')
parser.add_argument('--grid', type=float, default=(1.4, 1.2), nargs='+', metavar='1.4 1.2 1.2 1.3 ...',
                    help='Grid of alternating depth and width multipliers d and w (e.g.: d0, w0, d1, w1, ..., '
                         'dn, wn for n tuples)')
parser.add_argument('--phi', type=float, default=3.5, metavar=3.5,
                    help='Phi is the exponential scaling factor for width and depth multipliers (default: 3.5)')
parser.add_argument('--cuda', default=True, action='store_true',
                    help='By default CUDA training on GPU is enabled')
parser.add_argument('--no-cuda', dest='cuda', action='store_false',
                    help='Disables CUDA training and maps data plus model to CPU')
parser.add_argument('--dataset', default='CIFAR100', type=str, metavar='cifar',
                    help='Dataset to use. Choose from [CIFAR100, CIFAR10, MNIST, ImageNet] (default: CIFAR100)')
parser.add_argument('--image-size', default=32, type=int, metavar=32,
                    help='Input image size. Choose from [32 for CIFAR, 128-600 for ImageNet] (default: 32)')
parser.add_argument('--data-path', default='../data', type=str, metavar='/path',
                    help='Path to ImageNet data. CIFAR data will be downloaded to "../data" directory automatically '
                         'if the data-path argument is ignored')
parser.add_argument('--slurm-save', type=str, metavar='/path',
                    help='Path where slurm saves job results.')
parser.add_argument('--resume', default='./model_compound_scaling/saved_models/checkpoint.pt', type=str,
                    metavar='/path_cp',
                    help='Path to latest checkpoint. If it exists the train procedure will be resumed'
                         '(default: ./model_compound_scaling/saved_models/checkpoint.pt)')
parser.add_argument('--no-resume', dest='resume', action='store_false',
                    help='Do not resume from checkpoint')


def train(train_loader, model, criterion, optimizer, epoch):
    """
    Training procedure.

    Parameters:
    -----------
        train_loader:
            PyTorch Dataloader for given train dataset.
        model:
            The neural network model.
        criterion:
            Loss function to use (e.g. cross entropy, MSE, ...).
        optimizer:
            Updating model parameters with Gradient Descent plus Nesterov momentum and weight decay.
        epoch:
            Current training epoch.

    Returns:
    --------
        running loss
        running accuracy
    """

    # Resetting variables for next epoch
    running_loss = 0.0
    running_accuracy = 0.0
    n_steps = 0

    # Iterating over training dataloader
    for input, target in train_loader:

        # Measure data loading time
        end = time.time()

        # Resetting variables for calculating current batch accuracy
        correct = 0
        total = 0

        # Map data to GPU if available
        if use_cuda:
            input, target = input.cuda(), target.cuda(non_blocking=True)

        # Forward pass
        output = model(input)

        # Benchmark mode can allocate large memory blocks during the very first forward to test algorithms
        if epoch == 0:
            torch.cuda.empty_cache()

        # Compute batch_loss
        batch_loss = criterion(output, target)

        # Compute accuracy
        _, predicted = output.max(1)
        total += target.size(0)
        correct += predicted.eq(target).sum().item()
        batch_accuracy = 100. * correct / total

        # Compute gradient and do optimization step
        optimizer.zero_grad()
        batch_loss.backward()
        optimizer.step()

        # Summing up batch losses and accuracies over each step
        running_loss += batch_loss.float()
        running_accuracy += batch_accuracy
        n_steps += 1

        # Printing preliminary results
        if n_steps % 50 == 0:
            total_steps = int(len(train_loader.dataset) / train_loader.batch_size) + 1
            print('Epoch: {}, step: {}/{}, running_loss: {:.4f}, batch_acc: {:.4f}, running_acc: {:.4f}, '
                  'elapsed time: {:.2f} s, max_mem_alloc: {:.2f} GB, max_mem_cache {:.2f} GB'.format(
                epoch, n_steps, total_steps, running_loss / n_steps, batch_accuracy, running_accuracy / n_steps,
                                             (time.time() - end) * 100, torch.cuda.max_memory_allocated() / 1024 ** 3,
                                             torch.cuda.max_memory_cached() / 1024 ** 3))

        del output
        del batch_loss

    return running_loss / n_steps, running_accuracy / n_steps


def evaluate(model, loss, val_iterator):
    """
    Evaluating the model (no backward path / opimization step)

    Parameters:
    -----------
        model:
            The neural network model.
        loss:
            Loss function used for optimization (e.g. cross entropy, MSE, ...).
        val_iterator:
            PyTorch Dataloader for dataset which should be evaluated.

    Returns:
    --------
        Validation loss.
        Validation accuracy.
    """

    # Initializing parameters
    loss_value = 0.0
    accuracy = 0.0
    total_samples = 0

    with torch.no_grad():

        # Iterating over validation dataloader
        for data, labels in val_iterator:

            # Resetting variables for calculating current batch accuracy
            correct = 0
            total = 0

            # Map data to GPU if available
            if use_cuda:
                data = data.cuda()
                labels = labels.cuda(non_blocking=True)

            n_batch_samples = labels.size()[0]
            logits = model(data)

            # Compute batch loss
            batch_loss = loss(logits, labels)

            # Compute batch accuracy
            _, predicted = logits.max(1)
            total += labels.size(0)
            correct += predicted.eq(labels).sum().item()
            batch_accuracy = 100. * correct / total

            # Summing up batch losses and accuracies over each step
            loss_value += batch_loss.float() * n_batch_samples
            accuracy += batch_accuracy * n_batch_samples
            total_samples += n_batch_samples

            del logits
            del batch_loss

        return loss_value / total_samples, accuracy / total_samples


def grid_search(train_loader, val_loader, criterion, alpha, beta):
    """
    Builds the model with given scaling factors, sets up optimizer and learning rate schedulers plus executes
    training and evaluation of the model. A checkpoint is created each epoch. Also the best model will be saved.

    Parameters:
    -----------
        train_loader:
            PyTorch Dataloader for given train dataset.
        val_loader:
            PyTorch Dataloader for given validation dataset.
        criterion:
            Loss function to use (e.g. cross entropy, MSE, ...).
        alpha:
            Scaling factor for model depth.
        beta:
            Scaling factor for model width.
    """

    # Initializing training variables
    best_acc = 0
    all_losses = []

    # Initializing log file
    if args.slurm_save == None:
        logfile = open('./model_compound_scaling/logfiles/logfile.txt', 'a+')
    else:
        logfile = open(args.slurm_save + '/logfile.txt', 'a+')
    logfile.write('depth multiplier: {}, width multiplier: {}\n'.format(alpha, beta))

    # Building the model
    if args.dataset == 'CIFAR100':
        model = micronet(d_multiplier=alpha, w_multiplier=beta, num_classes=100)

    elif args.dataset == 'CIFAR10':
        model = micronet(d_multiplier=alpha, w_multiplier=beta, num_classes=10)

    elif args.dataset == 'ImageNet':
        model = image_micronet(d_multiplier=alpha, w_multiplier=beta)

    elif args.dataset == 'MNIST':
        model = lenet5(d_multiplier=alpha, w_multiplier=beta)

    # If multipile GPUs are used
    if use_cuda and torch.cuda.device_count() > 1:
        print("Let's use", torch.cuda.device_count(), "GPUs!")
        model = nn.DataParallel(model)

    # Transfers model to device (GPU/CPU). Device is globally initialized.
    model.to(device)

    # Defining the optimizer
    optimizer = torch.optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum,
                                weight_decay=args.weight_decay, nesterov=True)

    # KERAS like summary of the model architecture
    # summary(your_model, input_size=(channels, H, W), batch_size=-1, device="cuda")
    if use_cuda:
        if args.dataset == 'CIFAR100' or args.dataset == 'CIFAR10':
            #summary(model, (3, 32, 32), batch_size=args.batch_size)
            print(model)

        elif args.dataset == 'ImageNet':
            #summary(model, (3, args.image_size, args.image_size), batch_size=args.batch_size)
            print(model)

        elif args.dataset == 'MNIST':
            #summary(model, (3, args.image_size, args.image_size), batch_size=args.batch_size)
            print(model)

    # Optionally resume from a checkpoint
    if args.resume:
        if os.path.isfile(args.resume):
            print("=> loading checkpoint '{}'".format(args.resume))
            checkpoint = torch.load(args.resume)
            args.start_epoch = checkpoint['epoch']
            best_acc = checkpoint['acc']
            model.load_state_dict(checkpoint['state_dict'])
            optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
            load_last_epoch = checkpoint['epoch']-1
            print("=> loaded checkpoint '{}' (epoch {})".format(args.resume, checkpoint['epoch']))
        else:
            print("=> no checkpoint found at '{}'".format(args.resume))
            load_last_epoch = -1
    else:
        load_last_epoch = -1

    # Learning rate schedulers for cifar_micronet and imagenet_micronet
    lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer,
                                                           T_max = args.epochs,
                                                           eta_min = 0,
                                                           last_epoch = load_last_epoch)

    # START TRAINING
    start_time = time.time()
    model.train()

    for epoch in range(args.start_epoch, args.epochs):

        print('current lr {:.5e}'.format(optimizer.param_groups[0]['lr']))

        # Executing training process
        running_loss, running_accuracy = train(train_loader, model, criterion, optimizer, epoch)

        # Evaluation
        model.eval()
        val_loss, val_accuracy = evaluate(model, criterion, val_loader)

        # Logging the accuracies
        all_losses += [(epoch, running_loss, val_loss, running_accuracy, val_accuracy)]
        print('Epoch {0} running loss {1:.3f} val loss {2:.3f}  running acc {3:.3f} '
              'val acc{4:.3f}  time {5:.3f}'.format(*all_losses[-1], time.time() - start_time))
        logfile.write('Epoch {0} running loss {1:.3f} val loss {2:.3f}  running acc {3:.3f} '
              'val acc{4:.3f}  time {5:.3f}\n'.format(*all_losses[-1], time.time() - start_time))

        # Saving checkpoint
        if args.slurm_save == None:
            cp_savepath = './model_compound_scaling/saved_models/checkpoint.pt'
            best_savepath = './model_compound_scaling/saved_models/best_model.pt'
        else:
            cp_savepath = args.slurm_save + '/checkpoint.pt'
            best_savepath = args.slurm_save + '/best_model.pt'

        torch.save({
            'epoch': epoch,
            'state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'acc': val_accuracy,
            'lr': optimizer.param_groups[0]['lr']
        }, cp_savepath)

        # Make a lr scheduler step
        lr_scheduler.step()

        # Checking if current epoch yielded best validation accuracy
        is_best = val_accuracy > best_acc
        best_acc = max(val_accuracy, best_acc)

        # If so, saving best model state_dict
        if is_best and epoch > 0:
            torch.save(model.state_dict(), best_savepath)

        # Switch back to train mode
        model.train()
        start_time = time.time()


class Cutout(object):
    """
    Randomly mask out one or more patches from an image.

    Parameters:
    -----------
        n_holes (int):
            Number of patches to cut out of each image.
        length (int):
            The length (in pixels) of each square patch.
    """

    def __init__(self, n_holes, length):
        self.n_holes = n_holes
        self.length = length

    def __call__(self, img):
        """
        Parameters:
        -----------
            img (Tensor):
                Tensor image of size (C, H, W).
        Returns:
        --------
            Tensor:
                Image with n_holes of dimension length x length cut out of it.
        """
        h = img.size(1)
        w = img.size(2)

        mask = np.ones((h, w), np.float32)

        for n in range(self.n_holes):
            y = np.random.randint(h)
            x = np.random.randint(w)

            y1 = np.clip(y - self.length // 2, 0, h)
            y2 = np.clip(y + self.length // 2, 0, h)
            x1 = np.clip(x - self.length // 2, 0, w)
            x2 = np.clip(x + self.length // 2, 0, w)

            mask[y1: y2, x1: x2] = 0.

        mask = torch.from_numpy(mask)
        mask = mask.expand_as(img)
        img = img * mask

        return img


def main():
    """
    --------------------------------------------- MAIN --------------------------------------------------------
    Loads the data and executes the grid search on depth and width scaling factors.

    """
    # Manual seed for reproducibility
    torch.manual_seed(363636)

    # Global instances
    global args, use_cuda, device
    # Instantiating the parser
    args = parser.parse_args()
    # Global CUDA flag
    use_cuda = args.cuda and torch.cuda.is_available()
    # Defining device and device's map locationo
    device = torch.device("cuda" if use_cuda else "cpu")
    print('chosen device: ', device)

    # Defining loss function and printing CUDA information (if available)
    if use_cuda:
        print("PyTorch version: ")
        print(torch.__version__)
        print("CUDA Version: ")
        print(torch.version.cuda)
        print("cuDNN version is: ")
        print(cudnn.version())
        cudnn.benchmark = True
        criterion = nn.CrossEntropyLoss().cuda()
    else:
        criterion = nn.CrossEntropyLoss()

    # Dataloaders for CIFAR, ImageNet and MNIST
    if args.dataset == 'CIFAR100':

        normalize = transforms.Normalize(mean=[x / 255.0 for x in [125.3, 123.0, 113.9]],
                                         std=[x / 255.0 for x in [63.0, 62.1, 66.7]])

        kwargs = {'num_workers': args.workers, 'pin_memory': True} if use_cuda else {}

        train_loader = torch.utils.data.DataLoader(
            datasets.CIFAR100(root=args.data_path, train=True, transform=transforms.Compose([
                transforms.RandomHorizontalFlip(),
                transforms.RandomCrop(32, 4),
                transforms.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.3, hue=0.075),
                transforms.ToTensor(),
                normalize,
                Cutout(n_holes=1, length=16),
            ]), download=True),
            batch_size=args.batch_size, shuffle=True, **kwargs)

        val_loader = torch.utils.data.DataLoader(
            datasets.CIFAR100(root=args.data_path, train=False, transform=transforms.Compose([
                transforms.ToTensor(),
                normalize,
            ])),
            batch_size=args.val_batch_size, shuffle=False, **kwargs)

    elif args.dataset == 'ImageNet':

        traindir = os.path.join(args.data_path, 'train')
        valdir = os.path.join(args.data_path, 'val')
        normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                         std=[0.229, 0.224, 0.225])

        train_dataset = datasets.ImageFolder(
            traindir,
            transforms.Compose([
                transforms.RandomResizedCrop(args.image_size),
                transforms.RandomHorizontalFlip(),
                #transforms.RandomAffine(degrees=15),
                #transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.05),
                transforms.ToTensor(),
                normalize,
            ]))
        train_loader = torch.utils.data.DataLoader(
            train_dataset, batch_size=args.batch_size, shuffle=True,
            num_workers=args.workers, pin_memory=True)

        image_size = args.image_size
        val_dataset = datasets.ImageFolder(
            valdir,
            transforms.Compose([
                transforms.Resize(image_size, interpolation=PIL.Image.BICUBIC),
                transforms.CenterCrop(image_size),
                transforms.ToTensor(),
                normalize,
            ]))
        val_loader = torch.utils.data.DataLoader(
            val_dataset, batch_size=args.val_batch_size, shuffle=False,
            num_workers=args.workers, pin_memory=True)

    elif args.dataset == 'MNIST':

        kwargs = {'num_workers': args.workers, 'pin_memory': True} if use_cuda else {}

        train_loader = torch.utils.data.DataLoader(
            datasets.MNIST(args.data_path, train=True, download=True,
                           transform=transforms.Compose([
                               transforms.ToTensor(),
                               transforms.Normalize((0.1307,), (0.3081,))
                           ])),
            batch_size=args.batch_size, shuffle=True, **kwargs)
        val_loader = torch.utils.data.DataLoader(
            datasets.MNIST(args.data_path, train=False, transform=transforms.Compose([
                transforms.ToTensor(),
                transforms.Normalize((0.1307,), (0.3081,))
            ])),
            batch_size=args.val_batch_size, shuffle=True, **kwargs)

    elif args.dataset == 'CIFAR10':

        normalize = transforms.Normalize(mean=[x / 255.0 for x in [125.3, 123.0, 113.9]],
                                         std=[x / 255.0 for x in [63.0, 62.1, 66.7]])

        kwargs = {'num_workers': args.workers, 'pin_memory': True} if use_cuda else {}

        train_loader = torch.utils.data.DataLoader(
            datasets.CIFAR10(root=args.data_path, train=True, transform=transforms.Compose([
                transforms.RandomHorizontalFlip(),
                transforms.RandomCrop(32, 4),
                transforms.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.3, hue=0.075),
                transforms.ToTensor(),
                normalize,
                Cutout(n_holes=1, length=16),
            ]), download=True),
            batch_size=args.batch_size, shuffle=True, **kwargs)

        val_loader = torch.utils.data.DataLoader(
            datasets.CIFAR10(root=args.data_path, train=False, transform=transforms.Compose([
                transforms.ToTensor(),
                normalize,
            ])),
            batch_size=args.val_batch_size, shuffle=False, **kwargs)

    # original grid = [(1.0, 1.0), (1.9, 1.0), (1.7, 1.1), (1.6, 1.1), (1.4, 1.2), (1.2, 1.3), (1.0, 1.4)]

    grid = [(args.grid[i], args.grid[i+1]) for i in range(0, len(args.grid), 2)]

    for coeff in grid:
        alpha = coeff[0] ** args.phi
        beta = coeff[1] ** args.phi
        grid_search(train_loader, val_loader, criterion, alpha, beta)


if __name__ == '__main__':
    t1_start = time.perf_counter()
    main()
    t1_stop = time.perf_counter()
    print("Elapsed time: %.2f [s]" % (t1_stop - t1_start))

References:     https://github.com/TropComplique/trained-ternary-quantization     MIT License - Copyright (c) 2017 <NAME>     https://github.com/uoguelph-mlrg/Cutout     Educational Community License, Version 2.0 (ECL-2.0) - Copyright (c) 2019 <NAME>     https://github.com/lukemelas/EfficientNet-PyTorch     Apache License, Version 2.0 - Copyright (c) 2019 <NAME>     https://github.com/akamaster/pytorch_resnet_cifar10     <NAME>'s ResNet implementation for CIFAR10/CIFAR100 in PyTorch  This document is the main document for applying grid search on depth and width scaling factors for the MicroNet architectures which can be found in ./model/model.  This document contains the following functions/classes:      - main()         - Cutout(object)         - grid_search(train_loader, val_loader, criterion, alpha, beta)             - train(train_loader, model, criterion, optimizer, epoch)             - evaluate(model, loss, val_iterator) #ssfrom torchsummary import summary Training procedure.      Parameters:     -----------         train_loader:             PyTorch Dataloader for given train dataset.         model:             The neural network model.         criterion:             Loss function to use (e.g. cross entropy, MSE, ...).         optimizer:             Updating model parameters with Gradient Descent plus Nesterov momentum and weight decay.         epoch:             Current training epoch.      Returns:     --------         running loss         running accuracy # Resetting variables for next epoch # Iterating over training dataloader # Measure data loading time # Resetting variables for calculating current batch accuracy # Map data to GPU if available # Forward pass # Benchmark mode can allocate large memory blocks during the very first forward to test algorithms # Compute batch_loss # Compute accuracy # Compute gradient and do optimization step # Summing up batch losses and accuracies over each step # Printing preliminary results Evaluating the model (no backward path / opimization step)      Parameters:     -----------         model:             The neural network model.         loss:             Loss function used for optimization (e.g. cross entropy, MSE, ...).         val_iterator:             PyTorch Dataloader for dataset which should be evaluated.      Returns:     --------         Validation loss.         Validation accuracy. # Initializing parameters # Iterating over validation dataloader # Resetting variables for calculating current batch accuracy # Map data to GPU if available # Compute batch loss # Compute batch accuracy # Summing up batch losses and accuracies over each step Builds the model with given scaling factors, sets up optimizer and learning rate schedulers plus executes     training and evaluation of the model. A checkpoint is created each epoch. Also the best model will be saved.      Parameters:     -----------         train_loader:             PyTorch Dataloader for given train dataset.         val_loader:             PyTorch Dataloader for given validation dataset.         criterion:             Loss function to use (e.g. cross entropy, MSE, ...).         alpha:             Scaling factor for model depth.         beta:             Scaling factor for model width. # Initializing training variables # Initializing log file # Building the model # If multipile GPUs are used # Transfers model to device (GPU/CPU). Device is globally initialized. # Defining the optimizer # KERAS like summary of the model architecture # summary(your_model, input_size=(channels, H, W), batch_size=-1, device="cuda") #summary(model, (3, 32, 32), batch_size=args.batch_size) #summary(model, (3, args.image_size, args.image_size), batch_size=args.batch_size) #summary(model, (3, args.image_size, args.image_size), batch_size=args.batch_size) # Optionally resume from a checkpoint # Learning rate schedulers for cifar_micronet and imagenet_micronet # START TRAINING # Executing training process # Evaluation # Logging the accuracies # Saving checkpoint # Make a lr scheduler step # Checking if current epoch yielded best validation accuracy # If so, saving best model state_dict # Switch back to train mode Randomly mask out one or more patches from an image.      Parameters:     -----------         n_holes (int):             Number of patches to cut out of each image.         length (int):             The length (in pixels) of each square patch. Parameters:         -----------             img (Tensor):                 Tensor image of size (C, H, W).         Returns:         --------             Tensor:                 Image with n_holes of dimension length x length cut out of it. --------------------------------------------- MAIN --------------------------------------------------------     Loads the data and executes the grid search on depth and width scaling factors. # Manual seed for reproducibility # Global instances # Instantiating the parser # Global CUDA flag # Defining device and device's map locationo # Defining loss function and printing CUDA information (if available) # Dataloaders for CIFAR, ImageNet and MNIST #transforms.RandomAffine(degrees=15), #transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.05), # original grid = [(1.0, 1.0), (1.9, 1.0), (1.7, 1.1), (1.6, 1.1), (1.4, 1.2), (1.2, 1.3), (1.0, 1.4)]