DocumentAIKit/Dewarp/train_model.py

''' Document Localization using Recursive CNN
 Maintainer : Khurram Javed
 Email : kjaved@ualberta.ca '''

from __future__ import print_function

import argparse

import torch
import torch.utils.data as td

import dataprocessor
import experiment as ex
import model
import trainer
import utils

parser = argparse.ArgumentParser(description='Recursive-CNNs')
parser.add_argument('--batch-size', type=int, default=32, metavar='N',
                    help='input batch size for training (default: 32)')
parser.add_argument('--eval_interval', type=int, default=5)
parser.add_argument('--lr', type=float, default=0.005, metavar='LR',
                    help='learning rate (default: 0.005)')
parser.add_argument('--schedule', type=int, nargs='+', default=[10, 20, 30],
                    help='Decrease learning rate at these epochs.')
parser.add_argument('--gammas', type=float, nargs='+', default=[0.2, 0.2, 0.2],
                    help='LR is multiplied by gamma[k] on schedule[k], number of gammas should be equal to schedule')
parser.add_argument('--momentum', type=float, default=0.9, metavar='M',
                    help='SGD momentum (default: 0.9)')
parser.add_argument('--no-cuda', action='store_true', default=False,
                    help='disables CUDA training')
parser.add_argument('--pretrain', action='store_true', default=False,
                    help='Pretrain the model on CIFAR dataset?')
parser.add_argument('--load-ram', action='store_true', default=False,
                    help='Load data in ram: TODO : Remove this')
parser.add_argument('--debug', action='store_true', default=True,
                    help='Debug messages')
parser.add_argument('--seed', type=int, default=2323,
                    help='Seeds values to be used')
parser.add_argument('--log-interval', type=int, default=5, metavar='N',
                    help='how many batches to wait before logging training status')
parser.add_argument('--model-type', default="resnet",
                    help='model type to be used. Example : resnet32, resnet20, densenet, test')
parser.add_argument('--name', default="noname",
                    help='Name of the experiment')
parser.add_argument('--output-dir', default="./",
                    help='Directory to store the results; a new folder "DDMMYYYY" will be created '
                         'in the specified directory to save the results.')
parser.add_argument('--decay', type=float, default=0.00001, help='Weight decay (L2 penalty).')
parser.add_argument('--epochs', type=int, default=100, help='Number of epochs for trianing')
parser.add_argument('--dataset', default="document", help='Dataset to be used; example document, corner')
parser.add_argument('--loader', default="hdd", 
                    help='Loader to load data; hdd for reading from the hdd and ram for loading all data in the memory')
parser.add_argument("-i", "--data-dirs", nargs='+', default="/Users/khurramjaved96/documentTest64",
                    help="input Directory of train data")
parser.add_argument("-v", "--validation-dirs", nargs='+', default="/Users/khurramjaved96/documentTest64",
                    help="input Directory of val data")

args = parser.parse_args()

# Define an experiment.
my_experiment = ex.experiment(args.name, args, args.output_dir)

# Add logging support
logger = utils.utils.setup_logger(my_experiment.path)

args.cuda = not args.no_cuda and torch.cuda.is_available()

dataset = dataprocessor.DatasetFactory.get_dataset(args.data_dirs, args.dataset)

dataset_val = dataprocessor.DatasetFactory.get_dataset(args.validation_dirs, args.dataset)

# Fix the seed.
seed = args.seed
torch.manual_seed(seed)
if args.cuda:
    torch.cuda.manual_seed(seed)

train_dataset_loader = dataprocessor.LoaderFactory.get_loader(args.loader, dataset.myData,
                                                              transform=dataset.train_transform,
                                                              cuda=args.cuda)
# Loader used for training data
val_dataset_loader = dataprocessor.LoaderFactory.get_loader(args.loader, dataset_val.myData,
                                                            transform=dataset.test_transform,
                                                            cuda=args.cuda)
kwargs = {'num_workers': 1, 'pin_memory': True} if args.cuda else {}

# Iterator to iterate over training data.
train_iterator = torch.utils.data.DataLoader(train_dataset_loader,
                                             batch_size=args.batch_size, shuffle=True, **kwargs)
# Iterator to iterate over training data.
val_iterator = torch.utils.data.DataLoader(val_dataset_loader,
                                           batch_size=args.batch_size, shuffle=True, **kwargs)

# Get the required model
myModel = model.ModelFactory.get_model(args.model_type, args.dataset)
if args.cuda:
    myModel.cuda()

# Should I pretrain the model on CIFAR?
if args.pretrain:
    trainset = dataprocessor.DatasetFactory.get_dataset(None, "CIFAR")
    train_iterator_cifar = torch.utils.data.DataLoader(trainset, batch_size=32, shuffle=True, num_workers=2)

    # Define the optimizer used in the experiment
    cifar_optimizer = torch.optim.SGD(myModel.parameters(), args.lr, momentum=args.momentum,
                                      weight_decay=args.decay, nesterov=True)

    # Trainer object used for training
    cifar_trainer = trainer.CIFARTrainer(train_iterator_cifar, myModel, args.cuda, cifar_optimizer)

    for epoch in range(0, 70):
        logger.info("Epoch : %d", epoch)
        cifar_trainer.update_lr(epoch, [30, 45, 60], args.gammas)
        cifar_trainer.train(epoch)

    # Freeze the model
    counter = 0
    for name, param in myModel.named_parameters():
        # Getting the length of total layers so I can freeze x% of layers
        gen_len = sum(1 for _ in myModel.parameters())
        if counter < int(gen_len * 0.5):
            param.requires_grad = False
            logger.warning(name)
        else:
            logger.info(name)
        counter += 1

# Define the optimizer used in the experiment
optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad, myModel.parameters()), args.lr,
                            momentum=args.momentum,
                            weight_decay=args.decay, nesterov=True)

# Trainer object used for training
my_trainer = trainer.Trainer(train_iterator, myModel, args.cuda, optimizer)

# Evaluator
my_eval = trainer.EvaluatorFactory.get_evaluator("rmse", args.cuda)
# Running epochs_class epochs
cnt = 1
max_loss = 100
for epoch in range(0, args.epochs):
    logger.info("Epoch : %d", epoch)
    my_trainer.update_lr(epoch, args.schedule, args.gammas)
    l = my_trainer.train(epoch)
    loss = float(l)
    if loss < max_loss:
        torch.save(myModel.state_dict(), my_experiment.path + "_" + args.dataset + "_" + args.model_type+ "_best.pth")
        logger.info("best_model saved, avg_loss is %s", l)
        max_loss = loss
    logger.info("the best model's avg_loss is %s", l)
    if cnt%args.eval_interval == 0:
        my_eval.evaluate(my_trainer.model, val_iterator)
    cnt += 1
torch.save(myModel.state_dict(), my_experiment.path + "_" + args.dataset + "_" + args.model_type+ "_final.pth")
# torch.save(myModel, my_experiment.path + args.dataset + "_" + args.model_type+ ".pth")
my_experiment.store_json()