UPFlow_pytorch/scripts/simple_train.py

# -*- coding: utf-8 -*-
import os
from utils.tools import tools
import cv2
import numpy as np
from copy import deepcopy
import torch
import warnings  # ignore warnings
import torch.nn.functional as F
import torch.optim as optim
from dataset.kitti_dataset import kitti_train, kitti_flow
from model.upflow import UPFlow_net
from torch.utils.data import DataLoader
import time

''' scripts for training：
1. simply using photo loss and smooth loss
2. add occlusion checking
3. add teacher-student loss(ARFlow)
'''

# save and log loss value during training
class Loss_manager():
    def __init__(self):
        self.error_meter = tools.Avg_meter_ls()

    def fetch_loss(self, loss, loss_dict, name, batch_N, short_name=None):
        if name not in loss_dict.keys():
            pass
        elif loss_dict[name] is None:
            pass
        else:
            this_loss = loss_dict[name].mean()
            self.error_meter.update(name=name, val=this_loss.item(), num=batch_N, short_name=short_name)
            loss = loss + this_loss
        return loss

    def prepare_epoch(self):
        self.error_meter.reset()

    def log_info(self):
        p_str = self.error_meter.print_all_losses()
        return p_str

    def compute_loss(self, loss_dict, batch_N):
        loss = 0
        loss = self.fetch_loss(loss=loss, loss_dict=loss_dict, name='photo_loss', short_name='ph', batch_N=batch_N)
        loss = self.fetch_loss(loss=loss, loss_dict=loss_dict, name='smooth_loss', short_name='sm', batch_N=batch_N)
        loss = self.fetch_loss(loss=loss, loss_dict=loss_dict, name='census_loss', short_name='cen', batch_N=batch_N)
        # photo_loss, smooth_loss, census_loss = output_dict['photo_loss'].mean(), output_dict['smooth_loss'], output_dict['census_loss']
        loss = self.fetch_loss(loss=loss, loss_dict=loss_dict, name='msd_loss', short_name='msd', batch_N=batch_N)
        loss = self.fetch_loss(loss=loss, loss_dict=loss_dict, name='eq_loss', short_name='eq', batch_N=batch_N)
        loss = self.fetch_loss(loss=loss, loss_dict=loss_dict, name='oi_loss', short_name='oi', batch_N=batch_N)
        return loss

class Eval_model(tools.abs_test_model):
    def __init__(self):
        super(Eval_model, self).__init__()
        self.net_work = None

    def eval_forward(self, im1, im2, gt, *args):
        if self.net_work is None:
            raise ValueError('not network for evaluation')
        # === network output
        with torch.no_grad():
            input_dict = {'im1': im1, 'im2': im2, 'if_loss': False}
            output_dict = self.net_work(input_dict)
            flow_fw, flow_bw = output_dict['flow_f_out'], output_dict['flow_b_out']
            pred_flow = flow_fw
        return pred_flow

    def eval_save_result(self, save_name, predflow, *args, **kwargs):
        # you can save flow results here
        # print(save_name)
        pass

    def change_model(self, net):
        net.eval()
        self.net_work = net


class Trainer():
    class Config(tools.abstract_config):
        def __init__(self, **kwargs):
            self.exp_dir = './demo_exp'
            self.if_cuda = True

            self.batchsize = 2
            self.NUM_WORKERS = 4
            self.n_epoch = 1000
            self.batch_per_epoch = 500
            self.batch_per_print = 20
            self.lr = 1e-4
            self.weight_decay = 1e-4
            self.scheduler_gamma = 1

            # init
            self.update(kwargs)

        def __call__(self, ):
            t = Trainer(self)
            return t

    def __init__(self, conf: Config):
        self.conf = conf

        tools.check_dir(self.conf.exp_dir)

        # load network
        self.net = self.load_model()

        # for evaluation
        self.bench = self.load_eval_bench()
        self.eval_model = Eval_model()

        # load training dataset
        self.train_set = self.load_training_dataset()

    def training(self):
        train_loader = tools.data_prefetcher(self.train_set, batch_size=self.conf.batchsize, shuffle=True, num_workers=self.conf.NUM_WORKERS, pin_memory=True, drop_last=True)
        optimizer = optim.Adam(self.net.parameters(), lr=self.conf.lr, amsgrad=True, weight_decay=self.conf.weight_decay)
        scheduler = optim.lr_scheduler.ExponentialLR(optimizer, gamma=self.conf.scheduler_gamma)
        loss_manager = Loss_manager()
        timer = tools.time_clock()
        print("start training" + '=' * 10)
        i_batch = 0
        epoch = 0
        loss_manager.prepare_epoch()
        current_val, best_val, best_epoch = 0, 0, 0
        timer.start()
        while True:
            # prepare batch data
            batch_value = train_loader.next()
            if batch_value is None:
                batch_value = train_loader.next()
                assert batch_value is not None
            batchsize = batch_value['im1'].shape[0] # check if the im1 exists
            i_batch += 1
            # train batch
            self.net.train()
            optimizer.zero_grad()
            out_data = self.net(batch_value)  #

            loss_dict = out_data['loss_dict']
            loss = loss_manager.compute_loss(loss_dict=loss_dict, batch_N=batchsize)

            loss.backward()
            optimizer.step()
            if i_batch % self.conf.batch_per_print == 0:
                pass
            if i_batch % self.conf.batch_per_epoch == 0:
                # do eval  and check if save model todo===
                epoch+=1
                timer.end()
                print(' === epoch use time %.2f' % timer.get_during())
                scheduler.step(epoch=epoch)
                timer.start()

    def evaluation(self):
        self.eval_model.change_model(self.net)
        epe_all, f1, epe_noc, epe_occ = self.bench(self.eval_model)
        print('EPE All = %.2f, F1 = %.2f, EPE Noc = %.2f, EPE Occ = %.2f' % (epe_all, f1, epe_noc, epe_occ))
        print_str = 'EPE_%.2f__F1_%.2f__Noc_%.2f__Occ_%.2f' % (epe_all, f1, epe_noc, epe_occ)
        return epe_all, print_str

    # ======
    def load_model(self):
        param_dict = {
            # use cost volume norm
            'if_norm_before_cost_volume': True,
            'norm_moments_across_channels': False,
            'norm_moments_across_images': False,
            'if_froze_pwc': False,
            'if_use_cor_pytorch': False,  # speed is very slow, just for debug when cuda correlation is not compiled
            'if_sgu_upsample': False,  # 先把这个关掉跑通吧
        }
        pretrain_path = None  # pretrain path
        net_conf = UPFlow_net.config()
        net_conf.update(param_dict)
        net = net_conf()  # .cuda()
        if pretrain_path is not None:
            net.load_model(pretrain_path, if_relax=True, if_print=False)
        if self.conf.if_cuda:
            net = net.cuda()
        return net

    def load_eval_bench(self):
        bench = kitti_flow.Evaluation_bench(name='2015_train', if_gpu=self.conf.if_cuda, batch_size=1)
        return bench

    def load_training_dataset(self):
        data_config = {
            'crop_size': (256, 832),
            'rho': 8,
            'swap_images': True,
            'normalize': True,
            'horizontal_flip_aug': True,
        }
        data_conf = kitti_train.kitti_data_with_start_point.config(mv_type='2015', **data_config)
        dataset = data_conf()
        return dataset


if __name__ == '__main__':
    training_param = {}  # change param here
    conf = Trainer.Config(**training_param)
    trainer = conf()
    trainer.training()