train.py

import torch.optim as optim
import torch
from utils.train_utils import *
import logging
import math
import importlib
import datetime
from time import time
import random
import munch
import yaml
import os
import sys
import argparse
from dataset import ShapeNetH5


def train():
    logging.info(str(args))
    # Set GPU to use
    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
    os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
    
    metrics = ['cd_p', 'cd_t', 'emd', 'f1']
    best_epoch_losses = {m: (0, 0) if m == 'f1' else (0, math.inf) for m in metrics}
    #train_loss_meter = AverageValueMeter()
    val_loss_meters = {m: AverageValueMeter() for m in metrics}

    dataset = ShapeNetH5(train=True, npoints=args.num_points)
    dataset_test = ShapeNetH5(train=False, npoints=args.num_points)
    dataloader = torch.utils.data.DataLoader(dataset, batch_size=args.batch_size,
                                             shuffle=True, num_workers=int(args.workers))
    dataloader_test = torch.utils.data.DataLoader(dataset_test, batch_size=args.batch_size,
                                                  shuffle=False, num_workers=int(args.workers))
                                                  
    train_num = len(dataset) 
    test_num = len(dataset_test)    
    logging.info('Length of train dataset:%d', train_num)
    logging.info('Length of test dataset:%d', test_num)

    if not args.manual_seed:
        seed = random.randint(1, 10000)
    else:
        seed = int(args.manual_seed)
    logging.info('Random Seed: %d' % seed)
    random.seed(seed)
    torch.manual_seed(seed)

    model_module = importlib.import_module('.%s' % args.model_name, 'models')
    net = torch.nn.DataParallel(model_module.Model(args))
    net.cuda()
    if hasattr(model_module, 'weights_init'):
        net.module.apply(model_module.weights_init)

    cascade_gan = (args.model_name == 'cascade')
    net_d = None
    if cascade_gan:
        net_d = torch.nn.DataParallel(model_module.Discriminator(args))
        net_d.cuda()
        net_d.module.apply(model_module.weights_init)

    lr = args.lr
    if cascade_gan:
        lr_d = lr / 2
    if args.lr_decay:
        if args.lr_decay_interval and args.lr_step_decay_epochs:
            raise ValueError('lr_decay_interval and lr_step_decay_epochs are mutually exclusive!')
        if args.lr_step_decay_epochs:
            decay_epoch_list = [int(ep.strip()) for ep in args.lr_step_decay_epochs.split(',')]
            decay_rate_list = [float(rt.strip()) for rt in args.lr_step_decay_rates.split(',')]

    optimizer = getattr(optim, args.optimizer)
    if args.optimizer == 'Adagrad':
        optimizer = optimizer(net.module.parameters(), lr=lr, initial_accumulator_value=args.initial_accum_val)
    else:
        betas = args.betas.split(',')
        betas = (float(betas[0].strip()), float(betas[1].strip()))
        optimizer = optimizer(net.module.parameters(), lr=lr, weight_decay=args.weight_decay, betas=betas)

    if cascade_gan:
        optimizer_d = optim.Adam(net_d.parameters(), lr=lr_d, weight_decay=0.00001, betas=(0.5, 0.999))

    #scheduler = optim.lr_scheduler.OneCycleLR(optimizer,max_lr=lr*10, steps_per_epoch=len(dataloader), epochs=args.nepoch)
    scheduler = optim.lr_scheduler.OneCycleLR(optimizer, max_lr=lr*5, steps_per_epoch=len(dataloader), epochs=args.nepoch)

    alpha = None
    if args.varying_constant:
        varying_constant_epochs = [int(ep.strip()) for ep in args.varying_constant_epochs.split(',')]
        varying_constant = [float(c.strip()) for c in args.varying_constant.split(',')]
        assert len(varying_constant) == len(varying_constant_epochs) + 1

    if args.load_model:
        ckpt = torch.load(args.load_model)
        net.module.load_state_dict(ckpt['net_state_dict'])
        if cascade_gan:
            net_d.module.load_state_dict(ckpt['D_state_dict'])
        logging.info("%s's previous weights loaded." % args.model_name)

    for epoch in range(args.start_epoch + 1, args.nepoch + 1):
        #epoch_start_time = time()
        
        #train_loss_meter.reset()
        net.module.train()
        if cascade_gan:
            net_d.module.train()

        if args.varying_constant:
            for ind, ep in enumerate(varying_constant_epochs):
                if epoch < ep:
                    alpha = varying_constant[ind]
                    break
                elif ind == len(varying_constant_epochs)-1 and epoch >= ep:
                    alpha = varying_constant[ind+1]
                    break

        #if args.lr_decay:
        #    if args.lr_decay_interval:
        #        if epoch > 0 and epoch % args.lr_decay_interval == 0:
        #            lr = lr * args.lr_decay_rate
        #    elif args.lr_step_decay_epochs:
        #        if epoch in decay_epoch_list:
        #            lr = lr * decay_rate_list[decay_epoch_list.index(epoch)]
        #    if args.lr_clip:
        #        lr = max(lr, args.lr_clip)
        #    for param_group in optimizer.param_groups:
        #        param_group['lr'] = lr
        #        
        batch_end_time = time()
        for i, data in enumerate(dataloader, 0):
            data_time = time() - batch_end_time
        
            optimizer.zero_grad()
            if cascade_gan:
                optimizer_d.zero_grad()

            #_, inputs, gt = data
            inputs = data['partial_cloud']
            gt = data['gtcloud']

            inputs = inputs.float().cuda()
            gt = gt.float().cuda()
            inputs = inputs.transpose(2, 1).contiguous()
            # out2, loss2, net_loss = net(inputs, gt, mean_feature=mean_feature, alpha=alpha)
            out2, loss2, net_loss = net(inputs, gt, alpha=alpha)

            if cascade_gan:
                d_fake = generator_step(net_d, out2, net_loss, optimizer)
                discriminator_step(net_d, gt, d_fake, optimizer_d)
            else:
                #train_loss_meter.update(net_loss.mean().item())
                net_loss.backward(torch.squeeze(torch.ones(torch.cuda.device_count())).cuda())
                optimizer.step()
                scheduler.step()
                
            batch_time = time() - batch_end_time
            batch_end_time = time()
            if i % args.step_interval_to_print == 0:
                logging.info(' train [%d: %d/%d] BatchTime = %.3f(s) DataTime = %.3f(s), loss_type: %s, fine_loss: %f, total_loss: %f, lr: %f' %
                             (epoch, i, train_num / args.batch_size, batch_time, data_time, args.loss, loss2.mean().item(), net_loss.mean().item(), optimizer.param_groups[0]['lr']) + ' alpha: ' + str(alpha))

        if epoch % args.epoch_interval_to_save == 0:
            save_model('%s/network.pth' % log_dir, net, net_d=net_d)
            logging.info("Saving net...")

        if epoch % args.epoch_interval_to_val == 0 or epoch == args.nepoch - 1:
            val(net, epoch, val_loss_meters, dataloader_test, best_epoch_losses, test_num)


def val(net, curr_epoch_num, val_loss_meters, dataloader_test, best_epoch_losses, test_num):
    logging.info('Testing...')
    for v in val_loss_meters.values():
        v.reset()
    net.module.eval()

    with torch.no_grad():
        for i, data in enumerate(dataloader_test):
            #label, inputs, gt = data
            inputs = data['partial_cloud']
            gt = data['gtcloud']
            inputs = inputs.float().cuda()
            gt = gt.float().cuda()
            inputs = inputs.transpose(2, 1).contiguous()
            result_dict = net(inputs, gt, is_training=False)
            for k, v in val_loss_meters.items():
                v.update(result_dict[k].mean().item())

            logging.info(' test [%d: %d/%d], emd: %4f, cd_p: %4f, cd_t: %4f, f1: %f' %
                         (curr_epoch_num, i, test_num / args.batch_size, result_dict['emd'].mean().item()*1000,result_dict['cd_p'].mean().item()*1000, 
                          result_dict['cd_t'].mean().item()*1000, result_dict['f1'].mean().item()))
        fmt = 'best_%s: %f [epoch %d]; '
        best_log = ''
        for loss_type, (curr_best_epoch, curr_best_loss) in best_epoch_losses.items():
            if (val_loss_meters[loss_type].avg < curr_best_loss and loss_type != 'f1') or \
                    (val_loss_meters[loss_type].avg > curr_best_loss and loss_type == 'f1'):
                best_epoch_losses[loss_type] = (curr_epoch_num, val_loss_meters[loss_type].avg)
                save_model('%s/best_%s_network.pth' % (log_dir, loss_type), net)
                logging.info('Best %s net saved!' % loss_type)
                best_log += fmt % (loss_type, best_epoch_losses[loss_type][1], best_epoch_losses[loss_type][0])
            else:
                best_log += fmt % (loss_type, curr_best_loss, curr_best_epoch)

        curr_log = ''
        for loss_type, meter in val_loss_meters.items():
            curr_log += 'curr_%s: %f; ' % (loss_type, meter.avg)

        logging.info(curr_log)
        logging.info(best_log)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Train config file')
    parser.add_argument('--gpu', dest='gpu_id', help='GPU device to use', default='0,1', type=str)
    parser.add_argument('-c', '--config', help='path to config file', required=True)
    parse_args = parser.parse_args()
    config_path = parse_args.config
    args = munch.munchify(yaml.safe_load(open(config_path)))
    
    if parse_args.gpu_id is not None:
        args.gpu = parse_args.gpu_id

    log_time = datetime.datetime.now().isoformat()[:19]
    if args.load_model:
        exp_name = os.path.basename(os.path.dirname(args.load_model))
        log_dir = os.path.dirname(args.load_model)
    else:
        exp_name = args.model_name + '_' + args.loss + '_' + args.flag + '_' + log_time
        log_dir = os.path.join(args.work_dir, exp_name)
        if not os.path.exists(log_dir):
            os.makedirs(log_dir)
    logging.basicConfig(level=logging.INFO, handlers=[logging.FileHandler(os.path.join(log_dir, 'train.log')),
                                                     logging.StreamHandler(sys.stdout)])
    train()