train_model.py

#!/usr/bin/env python
# -*- coding:utf-8 _*-


import os
import torch
import torch.nn as nn
import torch.utils.data
from torch.utils.data import DataLoader
import logging
from tqdm import tqdm
from pointnet import PointNetCls
from Generator import Generator, Generator_origin
import numpy as np
from tensorboardX import SummaryWriter
import sklearn.metrics as metrics
import utils.data_utils as d_utils
import random
from utils import loss_utils
from dataloader import Modelnetload
from config import args

# Set random seed for reproducibility
manualSeed = 999
random.seed(manualSeed)
torch.manual_seed(manualSeed)


class Model:
    def __init__(self, args):
        self.opts = args
        self.backup()
        self.set_logger()
        self.Augmentor = args.Augmentor

    def backup(self):
        if not self.opts.restore:
            #source_folder = os.path.join(os.getcwd(), "Generate")
            common_folder = os.path.join(os.getcwd(), "utils")
            os.system("cp config.py '%s/model_cls.py.backup'" % (self.opts.log_dir))
            os.system("cp train_model.py '%s/train_model.py.backup'" % (self.opts.log_dir))
            os.system("cp Generator.py '%s/Generator.py.backup'" % (self.opts.log_dir))
            os.system("cp %s.py '%s/%s.py.backup'" % (self.opts.model_name, self.opts.log_dir, self.opts.model_name))
            os.system("cp %s/loss_utils.py '%s/loss_utils.py.backup'" % (common_folder, self.opts.log_dir))
            os.system("cp dataloader.py '%s/DataLoader.py.backup'" % (self.opts.log_dir))

    def set_logger(self):
        self.logger = logging.getLogger("CLS")
        self.logger.setLevel(logging.INFO)
        formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
        file_handler = logging.FileHandler(os.path.join(self.opts.log_dir, "log_train.txt"))
        file_handler.setLevel(logging.INFO)
        file_handler.setFormatter(formatter)
        self.logger.addHandler(file_handler)

    def train(self):

        self.log_string('PARAMETER ...')
        self.log_string(self.opts)
        with open(os.path.join(self.opts.log_dir, 'args.txt'), 'w') as log:
            for arg in sorted(vars(self.opts)):
                log.write(arg + ': ' + str(getattr(self.opts, arg)) + '\n')  # log of arguments
        writer = SummaryWriter(logdir=self.opts.log_dir)

        '''DATA LOADING'''
        self.log_string('Load dataset ...')

        trainDataLoader = DataLoader(Modelnetload(self.opts, partition='train'),
                                     batch_size=self.opts.batch_size, shuffle=True, drop_last=False)
        testDataLoader = DataLoader(Modelnetload(self.opts, partition='test'),
                                    batch_size=self.opts.batch_size, shuffle=False, )

        self.log_string("The number of training data is: %d" % len(trainDataLoader.dataset))
        self.log_string("The number of test data is: %d" % len(testDataLoader.dataset))

        '''MODEL LOADING'''
        num_class = 40
        use_cuda = torch.cuda.is_available()
        device = torch.device("cuda" if use_cuda else "cpu")
        self.dim = 3 if self.opts.use_normal else 0

        classifier = PointNetCls(num_class).cuda()
        if self.opts.generator is not None:
            generator = Generator(args).cuda()
        else:
            generator = Generator_origin().cuda()

        if torch.cuda.device_count() > 1:
            print("Let's use", torch.cuda.device_count(), "GPUs!")
            classifier = nn.DataParallel(classifier)
            generator = nn.DataParallel(generator)

        if self.opts.restore:
            self.log_string('Use pretrain Augment...')

            checkpoint = torch.load(self.opts.log_dir)
            start_epoch = checkpoint['epoch']
            classifier.load_state_dict(checkpoint['model_state_dict'])
        else:
            print('No existing Augment, starting training from scratch...')
            start_epoch = 0

        optimizer_c = torch.optim.Adam(
            classifier.parameters(),
            lr=self.opts.learning_rate,
            betas=(0.9, 0.999),
            eps=1e-08,
            weight_decay=self.opts.decay_rate
        )

        optimizer_a = torch.optim.Adam(
            generator.parameters(),
            lr=self.opts.learning_rate_a,
            betas=(0.9, 0.999),
            eps=1e-08,
            weight_decay=self.opts.decay_rate
        )
        if self.opts.no_decay:
            scheduler_c = None
        else:
            scheduler_c = torch.optim.lr_scheduler.StepLR(optimizer_c, step_size=20, gamma=self.opts.lr_decay)

        #scheduler_a = torch.optim.lr_scheduler.StepLR(optimizer_a, step_size=20, gamma=self.opts.lr_decay)
        scheduler_a = None

        global_epoch = 0
        best_tst_accuracy = 0.0
        blue = lambda x: '\033[94m' + x + '\033[0m'
        ispn = True if self.opts.model_name == "pointnet" else False

        '''TRAINING'''
        self.logger.info('Start training...')
        PointcloudScaleAndTranslate = d_utils.PointcloudScaleAndTranslate()  # initialize augmentation

        for epoch in range(start_epoch, self.opts.epoch):
            self.log_string('Epoch %d (%d/%s):' % (global_epoch + 1, epoch + 1, self.opts.epoch))
            if scheduler_c is not None:
                scheduler_c.step(epoch)
            if scheduler_a is not None:
                scheduler_a.step(epoch)

            for batch_id, data in tqdm(enumerate(trainDataLoader, 0), total=len(trainDataLoader), smoothing=0.9):
                if self.Augmentor is True:
                    origin, points, target = data
                    target = target[:, 0]
                    origin, points, target = origin.cuda(), points.cuda(), target.cuda().long()

                    origin = PointcloudScaleAndTranslate(origin)
                    origin = origin.transpose(2, 1).contiguous()
                    points = PointcloudScaleAndTranslate(points)
                    points = points.transpose(2, 1).contiguous()

                    noise = 0.02 * torch.randn(self.opts.batch_size, 1024).cuda()

                    classifier = classifier.train()
                    generator = generator.train()
                    optimizer_a.zero_grad()
                    optimizer_c.zero_grad()

                    gen_og = generator(origin, noise)
                    gen_pc = generator(points, noise)

                    pred_og, og_tran, og_feat = classifier(origin)
                    pred_gen, gen_tran, gen_feat = classifier(gen_og)
                    pred_pc, pc_tran, pc_feat = classifier(points)
                    pred_aug, aug_tran, aug_feat = classifier(gen_pc)

                    augLoss = loss_utils.aug_loss(pred_pc, pred_aug, pred_og, pred_gen, target,
                                                  pc_tran, aug_tran, og_tran, gen_tran, ispn=ispn)
                    clsLoss = loss_utils.cls_loss(pred_pc, pred_aug, pred_og, pred_gen, target,
                                                  pc_tran, aug_tran, og_tran, gen_tran,
                                                  pc_feat, aug_feat, og_feat, gen_feat, ispn=ispn)

                else:
                    points, target = data
                    target = target[:, 0]
                    points, target = points.cuda(), target.cuda().long()

                    points = PointcloudScaleAndTranslate(points)
                    points = points.transpose(2, 1).contiguous()

                    noise = 0.02 * torch.randn(self.opts.batch_size, 1024).cuda()

                    classifier = classifier.train()
                    generator = generator.train()
                    optimizer_a.zero_grad()
                    optimizer_c.zero_grad()

                    gen_pc = generator(points, noise)

                    pred_pc, pc_tran, pc_feat = classifier(points)
                    pred_aug, aug_tran, aug_feat = classifier(gen_pc)

                    augLoss = loss_utils.aug_loss_origin(pred_pc, pred_aug, target,
                                                        pc_tran, aug_tran, ispn=ispn)
                    clsLoss = loss_utils.cls_loss_origin(pred_pc, pred_aug, target,
                                                        pc_tran, aug_tran, pc_feat, aug_feat, ispn=ispn)

                augLoss.backward(retain_graph=True)
                clsLoss.backward(retain_graph=True)

                optimizer_a.step()
                optimizer_c.step()

            if self.Augmentor is True:
                train_acc = self.eval_one_epoch_train(classifier.eval(), trainDataLoader)
            else:
                train_acc = self.eval_one_epoch(classifier.eval(), trainDataLoader)
            test_acc = self.eval_one_epoch(classifier.eval(), testDataLoader)

            self.log_string('CLS Loss: %.2f' % clsLoss.data)
            self.log_string('AUG Loss: %.2f' % augLoss.data)

            self.log_string('Train Accuracy: %f' % train_acc)
            self.log_string('Test Accuracy: %f' % test_acc)

            writer.add_scalar("Train_Acc", train_acc, epoch)
            writer.add_scalar("Test_Acc", test_acc, epoch)

            if (test_acc >= best_tst_accuracy): #and test_acc >= 0.895:  # or (epoch % self.opts.epoch_per_save == 0):
                best_tst_accuracy = test_acc
                self.log_string('Save model...')
                self.save_checkpoint(
                    global_epoch + 1,
                    train_acc,
                    test_acc,
                    classifier,
                    optimizer_c,
                    str(self.opts.log_dir),
                    self.opts.model_name)

            global_epoch += 1
        self.log_string('Best Accuracy: %f' % best_tst_accuracy)
        self.log_string('End of training...')
        self.log_string(self.opts.log_dir)

    def eval_one_epoch(self, model, loader):
        mean_correct = []
        test_pred = []
        test_true = []

        for j, data in enumerate(loader, 0):

            points, target = data
            target = target[:, 0]
            points = points.transpose(2, 1)
            points, target = points.cuda(), target.cuda()
            classifier = model.eval()
            pred, _, _ = classifier(points)
            pred_choice = pred.data.max(1)[1]

            test_true.append(target.cpu().numpy())
            test_pred.append(pred_choice.detach().cpu().numpy())

            correct = pred_choice.eq(target.long().data).cpu().sum()
            mean_correct.append(correct.item() / float(points.size()[0]))

        test_true = np.concatenate(test_true)
        test_pred = np.concatenate(test_pred)
        test_acc = metrics.accuracy_score(test_true, test_pred)

        return test_acc

    def eval_one_epoch_train(self, model, loader):
        mean_correct = []
        test_pred = []
        test_true = []

        for j, data in enumerate(loader, 0):

            origin, point, target = data
            target = target[:, 0]
            points = origin.transpose(2, 1)
            points, target = points.cuda(), target.cuda()
            classifier = model.eval()
            pred, _, _ = classifier(points)
            pred_choice = pred.data.max(1)[1]

            test_true.append(target.cpu().numpy())
            test_pred.append(pred_choice.detach().cpu().numpy())

            correct = pred_choice.eq(target.long().data).cpu().sum()
            mean_correct.append(correct.item() / float(points.size()[0]))

        test_true = np.concatenate(test_true)
        test_pred = np.concatenate(test_pred)
        train_acc = metrics.accuracy_score(test_true, test_pred)

        return train_acc

    def save_checkpoint(self, epoch, train_accuracy, test_accuracy, model, optimizer, path, modelnet='checkpoint'):
        savepath = path + '/%s-%f-%04d.pth' % (modelnet, test_accuracy, epoch)
        print(savepath)
        state = {
            'epoch': epoch,
            'train_accuracy': train_accuracy,
            'test_accuracy': test_accuracy,
            # 'model_state_dict': model.module.state_dict(),
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
        }
        torch.save(state, savepath)

    def log_string(self, msg):
        print(msg)
        self.logger.info(msg)