<reponame>AkibMashrur/dicta-2021-evidential<gh_stars>0
# import general modules
import argparse
import os
import torch
import torch.backends.cudnn as cudnn

# import custom modules
import settings
import dataloader
import models
import utils

if __name__ == "__main__":
    # parsing input arguments
    parser = argparse.ArgumentParser(description='PyTorch training')
    parser.add_argument('--dataset', default='mnist',
                        help='Choose dataset mnist/cifar10')
    parser.add_argument('--model', default='cnn',
                        help='Choose model cnn/mcd/ensemble')
    parser.add_argument('--epochs', type=int, default=settings.EPOCHS,
                        help='Total number of epochs for training')
    parser.add_argument('--batch_size', type=int, default=settings.BATCH_SIZE,
                        help='Size of each training batch')
    parser.add_argument('--test_batch_size', type=int,
                        default=settings.TEST_BATCH_SIZE, help='Size of each testing batch')
    parser.add_argument('--seed', type=int,
                        default=settings.SEED, help='Random seed')
    parser.add_argument('--trials', type=int,
                        default=settings.TRIALS, help='Number of trials')
    args = parser.parse_args()

    # set root directory for saving model and logs
    ROOT_DIR_MODELS = './assets/saved_models/'
    ROOT_DIR_LOGS = './assets/saved_logs/'

    # set device
    device = 'cuda' if torch.cuda.is_available() else 'cpu'

    # ensure reporoducibility
    seed = args.seed

    # start experiments
    print(f"Starting trials for model {args.model} on {args.dataset}")
    for trial in range(args.trials):
        print(f"Running Trial #{trial+1}")
        # change seed at every trial
        seed = seed + trial
        torch.manual_seed(seed)
        if device == 'cuda':
            cudnn.benchmark = True
            torch.cuda.manual_seed(seed)

        # load model and store in available device
        model = models.get_model(args.model).to(device)
        train_loader, test_loader = dataloader.get_dataset(dataset_name=args.dataset,
                                                           batch_size=args.batch_size,
                                                           test_batch_size=args.test_batch_size)

        # start training
        epoch_train_losses, epoch_test_losses = [], []
        for epoch in range(args.epochs):
            train_loss = 0
            for batch_idx, (inputs, targets) in enumerate(train_loader):
                inputs, targets = inputs.to(device), targets.to(device)
                train_log = model.train_step(inputs, targets, epoch=epoch)
                train_loss = train_loss + train_log
            epoch_train_loss = train_loss/len(train_loader)
            epoch_train_losses.append(epoch_train_loss)

            test_loss = 0
            for batch_idx, (inputs, targets) in enumerate(test_loader):
                inputs, targets = inputs.to(device), targets.to(device)
                test_log = model.test_step(inputs, targets)
                test_loss = test_loss + test_log
            epoch_test_loss = (test_loss/len(test_loader))*100.
            epoch_test_losses.append(epoch_test_loss)

            print(
                f"Epoch {epoch+1}: Train loss: {epoch_train_loss:.3f}, Test accuracy: {epoch_test_loss:.3f}%")
        experiment_dir = ROOT_DIR_MODELS + 'T' + str(trial) + '_' + args.model  \
            + '_' + args.dataset + '_' + str(args.epochs) + 'e'
        experiment_dir_logs = ROOT_DIR_LOGS + 'T' + str(trial) + '_' + args.model  \
            + '_' + args.dataset + '_' + str(args.epochs) + 'e'

        # save trained model
        if not os.path.isdir(experiment_dir):
            os.makedirs(experiment_dir)

        torch.save(model.state_dict(), f'{experiment_dir}/final_model')

        # save model logs
        if not os.path.isdir(experiment_dir_logs):
            os.makedirs(experiment_dir_logs)
        utils.save_logs(epoch_train_losses, epoch_test_losses,
                        path=experiment_dir_logs)
        print(
            f"Trial #{trial+1} complete. Final train loss: {epoch_train_loss:.3f}, Final test accuracy: {epoch_test_loss:.3f}%")

    print("=== Training complete ===")
