# written by Jaekwang Cha
# version 0.1

# ================== IMPORT CUSTOM LEARNING LIBRARIES ===================== #
from customs.train import train, test
from customs.dataset import load_dataset
from customs.model import load_model

# ================== TRAINING SETTINGS ================== #
import argparse
import os

parser = argparse.ArgumentParser()
parser.add_argument('--train_method',    default='supervised',           type=str,      help='type of training: supervised(default), unsupervised, reinforce')
parser.add_argument('--task',            default='classification',       type=str,      help='task of training: classification(default), regression')
parser.add_argument('--dataset',         default='mnist',                type=str,      help='dataset to use')
parser.add_argument('--model',           default='CNN',                  type=str,      help='model to use')
parser.add_argument('--seed',            default=42,                     type=int,      help='random seed (default: 42)')
parser.add_argument('--num_worker',      default=1,                      type=int,      help='number of dataloader worker')
parser.add_argument('--no_cuda',         action='store_true',            default=False, help='disables CUDA training')
parser.add_argument('--gpu',             default=0,                      type=str,      help='GPU-id for GPU to use')
parser.add_argument('--multi_gpu',       default=0,                      type=str,      help='GPU-ids for multi-GPU usage')
parser.add_argument('--pin_memory',      default=True,                   type=bool,     help='pin memory option selector')
parser.add_argument('--save_model',      action='store_true',            default=False, help='For Saving the current Model')
parser.add_argument('--save_path',       default=os.getcwd()+'/weights', type=str,      help='Where to save weights')
parser.add_argument('--log_path',        default=os.getcwd()+'/Logs',    type=str,      help='Where to save Logs')


# data setting
parser.add_argument('--val_rate',        default=0.2,                    type=float,    help='split rate for the validation data')
parser.add_argument('--transform',       default='default',              type=str,      help='choose the data transform type')

# training parameter setting
parser.add_argument('--n_epoch',         default=10,                     type=int,      help='number of total training iteration')
parser.add_argument('--batch_size',      default=32,                     type=int,      help='size of minibatch')
parser.add_argument('--test_batch_size', default=32,                     type=int,      help='size of test-minibatch')

# optimizer & scheduler setting
parser.add_argument('--lr',              default=0.03,                   type=float,    help='training learning rate')
parser.add_argument('--optimizer',       default='adam',                 type=str,      help='optimizer select')
parser.add_argument('--scheduler',       default='steplr',               type=str,      help='scheduler select')


opt = parser.parse_args()


# ===================== IMPORT PYTORCH LIBRARIES ================== #
import torch
from torch.utils.data import DataLoader

torch.manual_seed(opt.seed)

# ================== GPU SETTINGS ================== #
def gpu_setup(opt):
    use_cuda = not opt.no_cuda and torch.cuda.is_available()
    os.environ["CUDA_DEVICE_ORDER"]   ="PCI_BUS_ID"

    if opt.multi_gpu != 0:
        print()
        print('Activating multi-gpu training mode')
        print(opt.multi_gpu)
        os.environ['CUDA_VISIBLE_DEVICES'] = str(opt.multi_gpu)
        opt.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    else:
        print()
        print('Activating single-gpu training mode')
        os.environ['CUDA_VISIBLE_DEVICES'] = str(opt.gpu)
        opt.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        print('Using gpu number ' + str(opt.gpu))
    return use_cuda

# ======================= MAIN SCRIPT ============================= #
def main(opt):
    use_cuda = gpu_setup(opt)
    dataset_train, dataset_validation = load_dataset(opt, train=True)
    print('training data size: {}'.format(len(dataset_train)))
    print('validation data size: {}'.format(len(dataset_validation)))

    dataset_test = load_dataset(opt, train=False)
    print('test data size: {}'.format(len(dataset_test)))
    print()
    
    kwargs = {'num_workers': opt.num_worker, 'pin_memory': opt.pin_memory} if use_cuda else {}
    train_dataloader = DataLoader(dataset_train, batch_size=opt.batch_size, shuffle=True, **kwargs)
    validation_dataloader = DataLoader(dataset_validation, batch_size=opt.batch_size, shuffle=True, **kwargs)
    test_dataloader = DataLoader(dataset_test, batch_size=opt.test_batch_size, shuffle=True, **kwargs)

    model = load_model(opt)
    if opt.multi_gpu != 0:
        model = torch.nn.DataParallel(model)
    model.to(opt.device)

    train(opt, model, train_dataloader, validation_dataloader)
    test(opt, model, test_dataloader)

if __name__ == '__main__':
    main(opt)

[((13, 9, 13, 34), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ({}, {}), '()', False, 'import argparse\n'), ((51, 0, 51, 27), 'torch.manual_seed', 'torch.manual_seed', ({(51, 18, 51, 26): 'opt.seed'}, {}), '(opt.seed)', False, 'import torch\n'), ((75, 40, 75, 69), 'customs.dataset.load_dataset', 'load_dataset', (), '', False, 'from customs.dataset import load_dataset\n'), ((79, 19, 79, 49), 'customs.dataset.load_dataset', 'load_dataset', (), '', False, 'from customs.dataset import load_dataset\n'), ((84, 23, 84, 99), 'torch.utils.data.DataLoader', 'DataLoader', (), '', False, 'from torch.utils.data import DataLoader\n'), ((85, 28, 85, 109), 'torch.utils.data.DataLoader', 'DataLoader', (), '', False, 'from torch.utils.data import DataLoader\n'), ((86, 22, 86, 102), 'torch.utils.data.DataLoader', 'DataLoader', (), '', False, 'from torch.utils.data import DataLoader\n'), ((88, 12, 88, 27), 'customs.model.load_model', 'load_model', ({(88, 23, 88, 26): 'opt'}, {}), '(opt)', False, 'from customs.model import load_model\n'), ((93, 4, 93, 62), 'customs.train.train', 'train', ({(93, 10, 93, 13): 'opt', (93, 15, 93, 20): 'model', (93, 22, 93, 38): 'train_dataloader', (93, 40, 93, 61): 'validation_dataloader'}, {}), '(opt, model, train_dataloader, validation_dataloader)', False, 'from customs.train import train, test\n'), ((94, 4, 94, 37), 'customs.train.test', 'test', ({(94, 9, 94, 12): 'opt', (94, 14, 94, 19): 'model', (94, 21, 94, 36): 'test_dataloader'}, {}), '(opt, model, test_dataloader)', False, 'from customs.train import train, test\n'), ((55, 35, 55, 60), 'torch.cuda.is_available', 'torch.cuda.is_available', ({}, {}), '()', False, 'import torch\n'), ((90, 16, 90, 44), 'torch.nn.DataParallel', 'torch.nn.DataParallel', ({(90, 38, 90, 43): 'model'}, {}), '(model)', False, 'import torch\n'), ((25, 49, 25, 60), 'os.getcwd', 'os.getcwd', ({}, {}), '()', False, 'import os\n'), ((26, 49, 26, 60), 'os.getcwd', 'os.getcwd', ({}, {}), '()', False, 'import os\n'), ((63, 44, 63, 69), 'torch.cuda.is_available', 'torch.cuda.is_available', ({}, {}), '()', False, 'import torch\n'), ((68, 44, 68, 69), 'torch.cuda.is_available', 'torch.cuda.is_available', ({}, {}), '()', False, 'import torch\n')]