<filename>train.py
import argparse
import os
import copy
import time
import random
import matplotlib.pyplot as plt
import numpy as np

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms

from models.parallel import DataParallelModel, DataParallelCriterion
from models.mobilenetv2 import MobileNetV2
from cfg.mobilenetv2_config import Config
from utils import CosineAnnealingWarmupRestarts, GradualWarmupScheduler

def set_random_seeds(random_seed, use_multi_gpu=False):
    '''Set random seeds.
    
    Args:
        random_seed (int): random seed to fix randomness
    '''
    torch.manual_seed(random_seed)
    torch.cuda.manual_seed(random_seed)
    if use_multi_gpu:
        torch.cuda.manual_seed_all(random_seed)

    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    np.random.seed(random_seed)
    random.seed(random_seed)

def load_data(data_dir, dataset='cifar10', batch_size=128):
    '''Load dataset.
    
    Args:
        data_dir (str): dataset path to load dataset
        dataset (int): dataset number (e.g. 0: CIFAR-10, 1: CIFAR-100)
    Returns:
        dataloaders (dict): data loaders for training and validation
        dataset_sizes (dict): each dataset size
    '''
    if 'cifar' in dataset:
        data_transforms = {
            'train': transforms.Compose([
                transforms.RandomCrop(32, padding=4),
                transforms.RandomHorizontalFlip(),
                transforms.ToTensor(),
                transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
            ]),
            'val': transforms.Compose([
                transforms.ToTensor(),
                transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
            ]),
        }
    elif dataset == 'imagenet':
        data_transforms = {
            'train': transforms.Compose([
                transforms.RandomResizedCrop(224),
                transforms.RandomHorizontalFlip(),
                transforms.ToTensor(),
                transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
            ]),
            'val': transforms.Compose([
                transforms.Resize(256),
                transforms.CenterCrop(224),
                transforms.ToTensor(),
                transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
            ]),
        }
    else:
        assert False, 'You choose wrong dataset.'

    # Download or load image datasets
    if dataset == 'cifar10':
        image_datasets = {x: datasets.CIFAR10(root='./data',
                                              train=True if x == 'train' else False,
                                              download=True, transform=data_transforms[x])
                             for x in ['train', 'val']}
    elif dataset == 'cifar100':
        image_datasets = {x: datasets.CIFAR100(root='./data',
                                               train=True if x == 'train' else False,
                                               download=True, transform=data_transforms[x])
                             for x in ['train', 'val']}
    elif dataset == 'imagenet':
        image_datasets = {x: datasets.ImageFolder(root=f'./data/imagenet/{x}',
                                                  transform=data_transforms[x])
                             for x in ['train', 'val']}
    else:
        assert False, 'Invalid Dataset.'

    # Create dataloaders
    dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=batch_size,
                                                  shuffle=True if x == 'train' else False,
                                                  num_workers=6)
                      for x in ['train', 'val']}
    dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']}    

    return dataloaders, dataset_sizes

def load_model(device, width_mult=1.0, use_res_connect=True, linear_bottleneck=True,
               res_loc=0, num_classes=10, pretrained_path='', use_multi_gpu=False,
               inverted_residual_setting=[], first_layer_stride=1, lr=0.1, t_max=200, eta_min=0):
    '''Load model, loss function, optimizer and scheduler.

    Args:
        device (obj): device object for cpu or gpu
        width_mult (float): multilplier value
        use_res_connect (bool): whether to use residual connection or not
        linear_bottleneck (bool): whether to use linear bottleneck or not
        res_loc (int): residual location
                       e.g. 0: between bottlenecks, 1: between expansions,
                            2: between depthwise layers
        num_classes (int): the number of classes
        pretrained_path (str): pretrained model path for transfer learning

    Returns:
        model (obj): loaded model
        criterion (obj): loss function
        optimizer (obj): optimizer (e.g. SGD, RMSprop)
        scheduler (obj): learning scheduler (e.g. Step, Cosine Annealing)
    '''
    model = MobileNetV2(num_classes=num_classes,
                        width_mult=width_mult,
                        use_res_connect=use_res_connect,
                        linear_bottleneck=linear_bottleneck,
                        res_loc=res_loc,
                        inverted_residual_setting=inverted_residual_setting,
                        first_layer_stride=first_layer_stride)

    # Device Settings (Single GPU or Multi-GPU)
    if use_multi_gpu:
        model = DataParallelModel(model).to(device)
        criterion = DataParallelCriterion(nn.CrossEntropyLoss())
    else:
        model = model.to(device)
        criterion = nn.CrossEntropyLoss()
    
    if pretrained_path:
        model.load_state_dict(torch.load(pretrained_path))

    # optimizer = optim.RMSprop(model.parameters(), lr=0.045, momentum=0.9, weight_decay=0.00004)
    optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9, weight_decay=4e-5)
    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=t_max, eta_min=eta_min)
    # scheduler = GradualWarmupScheduler(optimizer, multiplier=1, total_epoch=5, after_scheduler=scheduler)
    # scheduler = optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=10, T_mult=2)
    # scheduler = CosineAnnealingWarmupRestarts(optimizer,
    #                                           first_cycle_steps=150,
    #                                           cycle_mult=1.0,
    #                                           max_lr=lr,
    #                                           min_lr=0.0,
    #                                           warmup_steps=5,
    #                                           gamma=1.0)

    return model, criterion, optimizer, scheduler


def train_model(dataloaders, dataset_sizes, device,
                model, criterion, optimizer, scheduler,
                save_model, model_dir, save_acc, save_loss,
                result_dir, num_epochs=200, use_multi_gpu=False):
    '''Train and evaluate model.

    Args:
        dataloaders (dict): data loaders for training and validation
        dataset_sizes (dict): each dataset size
        device (obj): device object for cpu or gpu
        model (obj): loaded model
        criterion (obj): loss function
        optimizer (obj): optimizer (e.g. SGD, RMSprop)
        scheduler (obj): learning scheduler (e.g. Step, Cosine Annealing)
        save_model (bool) whether to save model or not
        model_dir (str): path to save model
        save_acc (bool): whether to save model accuracy or not
        result_dir (str): path to save accruacy
        num_epochs (int): the number of epochs

    Returns:
        model (obj): trained model
    '''
    since = time.time()

    best_model_wts = copy.deepcopy(model.state_dict())
    best_opt_wts = copy.deepcopy(optimizer.state_dict())
    best_acc = 0.0
    running_losses, running_accs = [], []
    id = round(time.time())

    for epoch in range(num_epochs):
        print('\nEpoch {}/{}'.format(epoch, num_epochs - 1))
        print('-' * 10)

        # Each epoch has a training and validation phase
        for phase in ['train', 'val']:
            if phase == 'train':
                model.train()  # Set model to training mode
            else:
                model.eval()   # Set model to evaluate mode

            running_loss = 0.0
            running_corrects = 0

            # Iterate over data.
            for inputs, labels in dataloaders[phase]:
                inputs = inputs.to(device)
                labels = labels.to(device)
                batch_size = inputs.size(0)

                # zero the parameter gradients
                optimizer.zero_grad()

                # forward
                # track history if only in train
                with torch.set_grad_enabled(phase == 'train'):
                    # Prediction
                    if use_multi_gpu:
                        outputs, gathered_outputs = model(inputs)
                        _, preds = torch.max(gathered_outputs, 1)
                    else:
                        outputs = model(inputs)
                        _, preds = torch.max(outputs, 1)

                    # Calculate a loss
                    loss = criterion(outputs, labels)

                    # backward + optimize only if in training phase
                    if phase == 'train':
                        loss.backward()
                        optimizer.step()

                # statistics
                running_loss += loss.item() * inputs.size(0)
                running_losses.append(loss.item())
                running_corrects += torch.sum(preds == labels.data)
                running_acc = torch.sum(preds == labels.data).double() / batch_size
                running_accs.append(running_acc.item())

            if phase == 'train':
                scheduler.step()

            epoch_loss = running_loss / dataset_sizes[phase]
            epoch_acc = running_corrects.double() / dataset_sizes[phase]

            print('{} Loss: {:.4f} Acc: {:.4f}'.format(
                phase, epoch_loss, epoch_acc))

            # deep copy the model
            if phase == 'val' and epoch_acc > best_acc:
                best_acc = epoch_acc
                best_model_wts = copy.deepcopy(model.state_dict())
                best_opt_wts = copy.deepcopy(optimizer.state_dict())

            if save_model:
                if not os.path.isdir(model_dir):
                    os.mkdir(model_dir)

                torch.save(model.state_dict(), f'{model_dir}/mobilenetv2_{id}_{epoch}')
                torch.save(best_opt_wts, f'{model_dir}/mobilenetv2_{id}_opt_{epoch}')

        print()

    time_elapsed = time.time() - since
    print('Training complete in {:.0f}m {:.0f}s'.format(
        time_elapsed // 60, time_elapsed % 60))
    print('Best val Acc: {:4f}'.format(best_acc))

    # load best model weights
    model.load_state_dict(best_model_wts)

    # save best model weights
    if save_model:
        if not os.path.isdir(model_dir):
            os.mkdir(model_dir)

        torch.save(model.state_dict(), f'{model_dir}/mobilenetv2_{id}_best')
        torch.save(best_opt_wts, f'{model_dir}/mobilenetv2_{id}_opt_best')

    # save accruacies
    # https://www.geeksforgeeks.org/graph-plotting-in-python-set-1/
    if save_acc:
        if not os.path.isdir(result_dir):
            os.mkdir(result_dir)

        steps = range(len(running_accs))
        
        # plotting the points  
        plt.plot(steps, running_accs)
        
        # naming the x and y axis 
        plt.xlabel('Steps') 
        plt.ylabel('Accruacy') 
        
        # giving a title to my graph 
        plt.title('Training Results') 

        plt.savefig(f'{result_dir}/acc_{id}.png')
        plt.clf() # https://www.activestate.com/resources/quick-reads/how-to-clear-a-plot-in-python/

        with open(f'{result_dir}/acc_{id}.txt', 'w') as f:
            f.write('\n'.join(list(map(str, running_accs))))

    if save_loss:
        # For Loss
        steps = range(len(running_losses))
        
        # plotting the points  
        plt.plot(steps, running_losses)
        
        # naming the x and y axis 
        plt.xlabel('Steps') 
        plt.ylabel('Loss') 
        
        # giving a title to my graph 
        plt.title('Training Results') 

        plt.savefig(f'{result_dir}/loss_{id}.png')

        with open(f'{result_dir}/loss_{id}.txt', 'w') as f:
            f.write('\n'.join(list(map(str, running_losses))))

    return model


if __name__ == '__main__':
    # Parse arguments
    parser = argparse.ArgumentParser(description='Train MobileNetV2')
    parser.add_argument('--data_dir', type=str, help='Path of input data', default='./data')
    parser.add_argument('--dataset', type=str, help='cifar10, cifar100, imagenet', default='cifar10')
    parser.add_argument('--save_model', action='store_true', help='Whether to save model or not')
    parser.add_argument('--pretrained_path', type=str, help='Path of model to save', default='')
    parser.add_argument('--model_dir', type=str, help='Path of model to save', default='./trained_models')
    parser.add_argument('--save_acc', action='store_true', help='Whether to save accruacies or not')
    parser.add_argument('--save_loss', action='store_true', help='Whether to save losses or not')
    parser.add_argument('--result_dir', type=str, help='Path of results to save', default='./results')
    parser.add_argument('--width_mult', type=float, help='Width for multiplier', default=1.0)
    parser.add_argument('--use_res_connect', action='store_true', help='Whether to use residual connection or not')
    parser.add_argument('--res_loc', type=int, help='Location of residual connections (e.g. Between bottlenecks -> 0', default=0)
    parser.add_argument('--linear_bottleneck', action='store_true', help='Whether to use Linear Bottlenck or not')
    parser.add_argument('--epoch', type=int, help='Epoch', default=200)
    parser.add_argument('--random_seed', type=int, help='Random seed for reproducibility', default=0)
    parser.add_argument('--batch_size', type=int, help='Batch Size', default=128)
    parser.add_argument('--print_to_file', action='store_true', help='Whether to print results or not')
    parser.add_argument('--use_multi_gpu', action='store_true', help='Whether to use multi-gpu or not')
    parser.add_argument('--lr', type=float, help='Learning rate', default=0.1)
    parser.add_argument('--t_max', type=int, help='T_max for cosine annealing', default=200)
    parser.add_argument('--eta_min', type=float, help='eta_min for cosine annealing', default=0.0)

    args = parser.parse_args()
    cfg = Config(args.dataset)

    # Set print function
    from utils import init_file_for_print, set_print_to_file
    id = round(time.time())
    init_file_for_print(id)
    print = set_print_to_file(print, args.print_to_file, id)

    # Set Random Seeds
    set_random_seeds(args.random_seed, args.use_multi_gpu)

    # Load dataset
    dataloaders, dataset_sizes = load_data(args.data_dir, args.dataset, args.batch_size)

    # Load MobileNetV2
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    model, criterion, optimizer, scheduler = load_model(device, args.width_mult,
                                                        args.use_res_connect,
                                                        args.linear_bottleneck,
                                                        args.res_loc,
                                                        cfg.num_cls,
                                                        args.pretrained_path,
                                                        args.use_multi_gpu,
                                                        cfg.inverted_residual_setting,
                                                        cfg.first_layer_stride,
                                                        args.lr, args.t_max, args.eta_min)
    
    # Train MobileNetV2
    model = train_model(dataloaders, dataset_sizes, device,
                        model, criterion, optimizer, scheduler,
                        args.save_model, args.model_dir,
                        args.save_acc, args.save_loss, args.result_dir,
                        args.epoch, args.use_multi_gpu)

<filename>train.py
import argparse
import os
import copy
import time
import random
import matplotlib.pyplot as plt
import numpy as np

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms

from models.parallel import DataParallelModel, DataParallelCriterion
from models.mobilenetv2 import MobileNetV2
from cfg.mobilenetv2_config import Config
from utils import CosineAnnealingWarmupRestarts, GradualWarmupScheduler

def set_random_seeds(random_seed, use_multi_gpu=False):
    '''Set random seeds.
    
    Args:
        random_seed (int): random seed to fix randomness
    '''
    torch.manual_seed(random_seed)
    torch.cuda.manual_seed(random_seed)
    if use_multi_gpu:
        torch.cuda.manual_seed_all(random_seed)

    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    np.random.seed(random_seed)
    random.seed(random_seed)

def load_data(data_dir, dataset='cifar10', batch_size=128):
    '''Load dataset.
    
    Args:
        data_dir (str): dataset path to load dataset
        dataset (int): dataset number (e.g. 0: CIFAR-10, 1: CIFAR-100)
    Returns:
        dataloaders (dict): data loaders for training and validation
        dataset_sizes (dict): each dataset size
    '''
    if 'cifar' in dataset:
        data_transforms = {
            'train': transforms.Compose([
                transforms.RandomCrop(32, padding=4),
                transforms.RandomHorizontalFlip(),
                transforms.ToTensor(),
                transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
            ]),
            'val': transforms.Compose([
                transforms.ToTensor(),
                transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
            ]),
        }
    elif dataset == 'imagenet':
        data_transforms = {
            'train': transforms.Compose([
                transforms.RandomResizedCrop(224),
                transforms.RandomHorizontalFlip(),
                transforms.ToTensor(),
                transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
            ]),
            'val': transforms.Compose([
                transforms.Resize(256),
                transforms.CenterCrop(224),
                transforms.ToTensor(),
                transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
            ]),
        }
    else:
        assert False, 'You choose wrong dataset.'

    # Download or load image datasets
    if dataset == 'cifar10':
        image_datasets = {x: datasets.CIFAR10(root='./data',
                                              train=True if x == 'train' else False,
                                              download=True, transform=data_transforms[x])
                             for x in ['train', 'val']}
    elif dataset == 'cifar100':
        image_datasets = {x: datasets.CIFAR100(root='./data',
                                               train=True if x == 'train' else False,
                                               download=True, transform=data_transforms[x])
                             for x in ['train', 'val']}
    elif dataset == 'imagenet':
        image_datasets = {x: datasets.ImageFolder(root=f'./data/imagenet/{x}',
                                                  transform=data_transforms[x])
                             for x in ['train', 'val']}
    else:
        assert False, 'Invalid Dataset.'

    # Create dataloaders
    dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=batch_size,
                                                  shuffle=True if x == 'train' else False,
                                                  num_workers=6)
                      for x in ['train', 'val']}
    dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']}    

    return dataloaders, dataset_sizes

def load_model(device, width_mult=1.0, use_res_connect=True, linear_bottleneck=True,
               res_loc=0, num_classes=10, pretrained_path='', use_multi_gpu=False,
               inverted_residual_setting=[], first_layer_stride=1, lr=0.1, t_max=200, eta_min=0):
    '''Load model, loss function, optimizer and scheduler.

    Args:
        device (obj): device object for cpu or gpu
        width_mult (float): multilplier value
        use_res_connect (bool): whether to use residual connection or not
        linear_bottleneck (bool): whether to use linear bottleneck or not
        res_loc (int): residual location
                       e.g. 0: between bottlenecks, 1: between expansions,
                            2: between depthwise layers
        num_classes (int): the number of classes
        pretrained_path (str): pretrained model path for transfer learning

    Returns:
        model (obj): loaded model
        criterion (obj): loss function
        optimizer (obj): optimizer (e.g. SGD, RMSprop)
        scheduler (obj): learning scheduler (e.g. Step, Cosine Annealing)
    '''
    model = MobileNetV2(num_classes=num_classes,
                        width_mult=width_mult,
                        use_res_connect=use_res_connect,
                        linear_bottleneck=linear_bottleneck,
                        res_loc=res_loc,
                        inverted_residual_setting=inverted_residual_setting,
                        first_layer_stride=first_layer_stride)

    # Device Settings (Single GPU or Multi-GPU)
    if use_multi_gpu:
        model = DataParallelModel(model).to(device)
        criterion = DataParallelCriterion(nn.CrossEntropyLoss())
    else:
        model = model.to(device)
        criterion = nn.CrossEntropyLoss()
    
    if pretrained_path:
        model.load_state_dict(torch.load(pretrained_path))

    # optimizer = optim.RMSprop(model.parameters(), lr=0.045, momentum=0.9, weight_decay=0.00004)
    optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9, weight_decay=4e-5)
    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=t_max, eta_min=eta_min)
    # scheduler = GradualWarmupScheduler(optimizer, multiplier=1, total_epoch=5, after_scheduler=scheduler)
    # scheduler = optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=10, T_mult=2)
    # scheduler = CosineAnnealingWarmupRestarts(optimizer,
    #                                           first_cycle_steps=150,
    #                                           cycle_mult=1.0,
    #                                           max_lr=lr,
    #                                           min_lr=0.0,
    #                                           warmup_steps=5,
    #                                           gamma=1.0)

    return model, criterion, optimizer, scheduler


def train_model(dataloaders, dataset_sizes, device,
                model, criterion, optimizer, scheduler,
                save_model, model_dir, save_acc, save_loss,
                result_dir, num_epochs=200, use_multi_gpu=False):
    '''Train and evaluate model.

    Args:
        dataloaders (dict): data loaders for training and validation
        dataset_sizes (dict): each dataset size
        device (obj): device object for cpu or gpu
        model (obj): loaded model
        criterion (obj): loss function
        optimizer (obj): optimizer (e.g. SGD, RMSprop)
        scheduler (obj): learning scheduler (e.g. Step, Cosine Annealing)
        save_model (bool) whether to save model or not
        model_dir (str): path to save model
        save_acc (bool): whether to save model accuracy or not
        result_dir (str): path to save accruacy
        num_epochs (int): the number of epochs

    Returns:
        model (obj): trained model
    '''
    since = time.time()

    best_model_wts = copy.deepcopy(model.state_dict())
    best_opt_wts = copy.deepcopy(optimizer.state_dict())
    best_acc = 0.0
    running_losses, running_accs = [], []
    id = round(time.time())

    for epoch in range(num_epochs):
        print('\nEpoch {}/{}'.format(epoch, num_epochs - 1))
        print('-' * 10)

        # Each epoch has a training and validation phase
        for phase in ['train', 'val']:
            if phase == 'train':
                model.train()  # Set model to training mode
            else:
                model.eval()   # Set model to evaluate mode

            running_loss = 0.0
            running_corrects = 0

            # Iterate over data.
            for inputs, labels in dataloaders[phase]:
                inputs = inputs.to(device)
                labels = labels.to(device)
                batch_size = inputs.size(0)

                # zero the parameter gradients
                optimizer.zero_grad()

                # forward
                # track history if only in train
                with torch.set_grad_enabled(phase == 'train'):
                    # Prediction
                    if use_multi_gpu:
                        outputs, gathered_outputs = model(inputs)
                        _, preds = torch.max(gathered_outputs, 1)
                    else:
                        outputs = model(inputs)
                        _, preds = torch.max(outputs, 1)

                    # Calculate a loss
                    loss = criterion(outputs, labels)

                    # backward + optimize only if in training phase
                    if phase == 'train':
                        loss.backward()
                        optimizer.step()

                # statistics
                running_loss += loss.item() * inputs.size(0)
                running_losses.append(loss.item())
                running_corrects += torch.sum(preds == labels.data)
                running_acc = torch.sum(preds == labels.data).double() / batch_size
                running_accs.append(running_acc.item())

            if phase == 'train':
                scheduler.step()

            epoch_loss = running_loss / dataset_sizes[phase]
            epoch_acc = running_corrects.double() / dataset_sizes[phase]

            print('{} Loss: {:.4f} Acc: {:.4f}'.format(
                phase, epoch_loss, epoch_acc))

            # deep copy the model
            if phase == 'val' and epoch_acc > best_acc:
                best_acc = epoch_acc
                best_model_wts = copy.deepcopy(model.state_dict())
                best_opt_wts = copy.deepcopy(optimizer.state_dict())

            if save_model:
                if not os.path.isdir(model_dir):
                    os.mkdir(model_dir)

                torch.save(model.state_dict(), f'{model_dir}/mobilenetv2_{id}_{epoch}')
                torch.save(best_opt_wts, f'{model_dir}/mobilenetv2_{id}_opt_{epoch}')

        print()

    time_elapsed = time.time() - since
    print('Training complete in {:.0f}m {:.0f}s'.format(
        time_elapsed // 60, time_elapsed % 60))
    print('Best val Acc: {:4f}'.format(best_acc))

    # load best model weights
    model.load_state_dict(best_model_wts)

    # save best model weights
    if save_model:
        if not os.path.isdir(model_dir):
            os.mkdir(model_dir)

        torch.save(model.state_dict(), f'{model_dir}/mobilenetv2_{id}_best')
        torch.save(best_opt_wts, f'{model_dir}/mobilenetv2_{id}_opt_best')

    # save accruacies
    # https://www.geeksforgeeks.org/graph-plotting-in-python-set-1/
    if save_acc:
        if not os.path.isdir(result_dir):
            os.mkdir(result_dir)

        steps = range(len(running_accs))
        
        # plotting the points  
        plt.plot(steps, running_accs)
        
        # naming the x and y axis 
        plt.xlabel('Steps') 
        plt.ylabel('Accruacy') 
        
        # giving a title to my graph 
        plt.title('Training Results') 

        plt.savefig(f'{result_dir}/acc_{id}.png')
        plt.clf() # https://www.activestate.com/resources/quick-reads/how-to-clear-a-plot-in-python/

        with open(f'{result_dir}/acc_{id}.txt', 'w') as f:
            f.write('\n'.join(list(map(str, running_accs))))

    if save_loss:
        # For Loss
        steps = range(len(running_losses))
        
        # plotting the points  
        plt.plot(steps, running_losses)
        
        # naming the x and y axis 
        plt.xlabel('Steps') 
        plt.ylabel('Loss') 
        
        # giving a title to my graph 
        plt.title('Training Results') 

        plt.savefig(f'{result_dir}/loss_{id}.png')

        with open(f'{result_dir}/loss_{id}.txt', 'w') as f:
            f.write('\n'.join(list(map(str, running_losses))))

    return model


if __name__ == '__main__':
    # Parse arguments
    parser = argparse.ArgumentParser(description='Train MobileNetV2')
    parser.add_argument('--data_dir', type=str, help='Path of input data', default='./data')
    parser.add_argument('--dataset', type=str, help='cifar10, cifar100, imagenet', default='cifar10')
    parser.add_argument('--save_model', action='store_true', help='Whether to save model or not')
    parser.add_argument('--pretrained_path', type=str, help='Path of model to save', default='')
    parser.add_argument('--model_dir', type=str, help='Path of model to save', default='./trained_models')
    parser.add_argument('--save_acc', action='store_true', help='Whether to save accruacies or not')
    parser.add_argument('--save_loss', action='store_true', help='Whether to save losses or not')
    parser.add_argument('--result_dir', type=str, help='Path of results to save', default='./results')
    parser.add_argument('--width_mult', type=float, help='Width for multiplier', default=1.0)
    parser.add_argument('--use_res_connect', action='store_true', help='Whether to use residual connection or not')
    parser.add_argument('--res_loc', type=int, help='Location of residual connections (e.g. Between bottlenecks -> 0', default=0)
    parser.add_argument('--linear_bottleneck', action='store_true', help='Whether to use Linear Bottlenck or not')
    parser.add_argument('--epoch', type=int, help='Epoch', default=200)
    parser.add_argument('--random_seed', type=int, help='Random seed for reproducibility', default=0)
    parser.add_argument('--batch_size', type=int, help='Batch Size', default=128)
    parser.add_argument('--print_to_file', action='store_true', help='Whether to print results or not')
    parser.add_argument('--use_multi_gpu', action='store_true', help='Whether to use multi-gpu or not')
    parser.add_argument('--lr', type=float, help='Learning rate', default=0.1)
    parser.add_argument('--t_max', type=int, help='T_max for cosine annealing', default=200)
    parser.add_argument('--eta_min', type=float, help='eta_min for cosine annealing', default=0.0)

    args = parser.parse_args()
    cfg = Config(args.dataset)

    # Set print function
    from utils import init_file_for_print, set_print_to_file
    id = round(time.time())
    init_file_for_print(id)
    print = set_print_to_file(print, args.print_to_file, id)

    # Set Random Seeds
    set_random_seeds(args.random_seed, args.use_multi_gpu)

    # Load dataset
    dataloaders, dataset_sizes = load_data(args.data_dir, args.dataset, args.batch_size)

    # Load MobileNetV2
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    model, criterion, optimizer, scheduler = load_model(device, args.width_mult,
                                                        args.use_res_connect,
                                                        args.linear_bottleneck,
                                                        args.res_loc,
                                                        cfg.num_cls,
                                                        args.pretrained_path,
                                                        args.use_multi_gpu,
                                                        cfg.inverted_residual_setting,
                                                        cfg.first_layer_stride,
                                                        args.lr, args.t_max, args.eta_min)
    
    # Train MobileNetV2
    model = train_model(dataloaders, dataset_sizes, device,
                        model, criterion, optimizer, scheduler,
                        args.save_model, args.model_dir,
                        args.save_acc, args.save_loss, args.result_dir,
                        args.epoch, args.use_multi_gpu)

Set random seeds.          Args:         random_seed (int): random seed to fix randomness Load dataset.          Args:         data_dir (str): dataset path to load dataset         dataset (int): dataset number (e.g. 0: CIFAR-10, 1: CIFAR-100)     Returns:         dataloaders (dict): data loaders for training and validation         dataset_sizes (dict): each dataset size # Download or load image datasets # Create dataloaders Load model, loss function, optimizer and scheduler.      Args:         device (obj): device object for cpu or gpu         width_mult (float): multilplier value         use_res_connect (bool): whether to use residual connection or not         linear_bottleneck (bool): whether to use linear bottleneck or not         res_loc (int): residual location                        e.g. 0: between bottlenecks, 1: between expansions,                             2: between depthwise layers         num_classes (int): the number of classes         pretrained_path (str): pretrained model path for transfer learning      Returns:         model (obj): loaded model         criterion (obj): loss function         optimizer (obj): optimizer (e.g. SGD, RMSprop)         scheduler (obj): learning scheduler (e.g. Step, Cosine Annealing) # Device Settings (Single GPU or Multi-GPU) # optimizer = optim.RMSprop(model.parameters(), lr=0.045, momentum=0.9, weight_decay=0.00004) # scheduler = GradualWarmupScheduler(optimizer, multiplier=1, total_epoch=5, after_scheduler=scheduler) # scheduler = optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=10, T_mult=2) # scheduler = CosineAnnealingWarmupRestarts(optimizer, #                                           first_cycle_steps=150, #                                           cycle_mult=1.0, #                                           max_lr=lr, #                                           min_lr=0.0, #                                           warmup_steps=5, #                                           gamma=1.0) Train and evaluate model.      Args:         dataloaders (dict): data loaders for training and validation         dataset_sizes (dict): each dataset size         device (obj): device object for cpu or gpu         model (obj): loaded model         criterion (obj): loss function         optimizer (obj): optimizer (e.g. SGD, RMSprop)         scheduler (obj): learning scheduler (e.g. Step, Cosine Annealing)         save_model (bool) whether to save model or not         model_dir (str): path to save model         save_acc (bool): whether to save model accuracy or not         result_dir (str): path to save accruacy         num_epochs (int): the number of epochs      Returns:         model (obj): trained model # Each epoch has a training and validation phase # Set model to training mode # Set model to evaluate mode # Iterate over data. # zero the parameter gradients # forward # track history if only in train # Prediction # Calculate a loss # backward + optimize only if in training phase # statistics # deep copy the model # load best model weights # save best model weights # save accruacies # https://www.geeksforgeeks.org/graph-plotting-in-python-set-1/ # plotting the points # naming the x and y axis # giving a title to my graph # https://www.activestate.com/resources/quick-reads/how-to-clear-a-plot-in-python/ # For Loss # plotting the points # naming the x and y axis # giving a title to my graph # Parse arguments # Set print function # Set Random Seeds # Load dataset # Load MobileNetV2 # Train MobileNetV2