model/depth_engine.py

import numpy as np
import os, time, datetime
import torch
import torch.utils.tensorboard
import importlib
import shutil
import utils.util as util
import utils.util_vis as util_vis

from torch.nn.parallel import DistributedDataParallel as DDP
from utils.util import print_eval, setup, cleanup
from utils.util import EasyDict as edict
from utils.eval_depth import DepthMetric
from copy import deepcopy
from model.compute_graph import graph_depth

# ============================ main engine for training and evaluation ============================

class Runner():

    def __init__(self, opt):
        super().__init__()
        if os.path.isdir(opt.output_path) and opt.resume == False and opt.device == 0:
            for filename in os.listdir(opt.output_path):
                if "tfevents" in filename: os.remove(os.path.join(opt.output_path, filename))
                if "html" in filename: os.remove(os.path.join(opt.output_path, filename))
                if "vis" in filename: shutil.rmtree(os.path.join(opt.output_path, filename))
                if "dump" in filename: shutil.rmtree(os.path.join(opt.output_path, filename))
                if "embedding" in filename: shutil.rmtree(os.path.join(opt.output_path, filename))
        if opt.device == 0: 
            os.makedirs(opt.output_path,exist_ok=True)
        setup(opt.device, opt.world_size, opt.port)
        opt.batch_size = opt.batch_size // opt.world_size
    
    def get_viz_data(self, opt):
        # get data for visualization
        viz_data_list = []
        sample_range = len(self.viz_loader)
        viz_interval = sample_range // opt.eval.n_vis
        for i in range(sample_range):
            current_batch = next(self.viz_loader_iter)
            if i % viz_interval != 0: continue
            viz_data_list.append(current_batch)
        return viz_data_list
    
    def load_dataset(self, opt, eval_split="test"):
        data_train = importlib.import_module('data.{}'.format(opt.data.dataset_train))
        data_test = importlib.import_module('data.{}'.format(opt.data.dataset_test))
        if opt.device == 0: print("loading training data...")
        self.batch_order = []
        self.train_data = data_train.Dataset(opt, split="train", load_3D=False)
        self.train_loader = self.train_data.setup_loader(opt, shuffle=True, use_ddp=True, drop_last=True)
        self.num_batches = len(self.train_loader)
        if opt.device == 0: print("loading test data...")
        self.test_data = data_test.Dataset(opt, split=eval_split, load_3D=False)
        self.test_loader = self.test_data.setup_loader(opt, shuffle=False, use_ddp=True, drop_last=True, batch_size=opt.eval.batch_size)
        self.num_batches_test = len(self.test_loader)
        if len(self.test_loader.sampler) * opt.world_size < len(self.test_data):
            self.aux_test_dataset = torch.utils.data.Subset(self.test_data,
                                    range(len(self.test_loader.sampler) * opt.world_size, len(self.test_data)))
            self.aux_test_loader = torch.utils.data.DataLoader(
                self.aux_test_dataset, batch_size=opt.eval.batch_size, shuffle=False, drop_last=False,
                num_workers=opt.data.num_workers)
        if opt.device == 0:
            print("creating data for visualization...")
            self.viz_loader = self.test_data.setup_loader(opt, shuffle=False, use_ddp=False, drop_last=False, batch_size=1)
            self.viz_loader_iter = iter(self.viz_loader)
            self.viz_data = self.get_viz_data(opt)

    def build_networks(self, opt):
        if opt.device == 0: print("building networks...")
        self.graph = DDP(graph_depth.Graph(opt).to(opt.device), device_ids=[opt.device], find_unused_parameters=True)
        self.depth_metric = DepthMetric(thresholds=opt.eval.d_thresholds, depth_cap=opt.eval.depth_cap)

# =================================================== set up training =========================================================

    def setup_optimizer(self, opt):
        if opt.device == 0: print("setting up optimizers...")
        param_nodecay = []
        param_decay = []
        for name, param in self.graph.named_parameters():
            # skip and fixed params
            if not param.requires_grad:
                continue
            if param.ndim <= 1 or name.endswith(".bias"):
                # print("{} -> finetune_param_nodecay".format(name))
                param_nodecay.append(param)
            else:
                param_decay.append(param)
                # print("{} -> finetune_param_decay".format(name))
            # create the optim dictionary
            optim_dict = [
                {'params': param_nodecay, 'lr': opt.optim.lr, 'weight_decay': 0.},
                {'params': param_decay, 'lr': opt.optim.lr, 'weight_decay': opt.optim.weight_decay}
            ]
            
        self.optim = torch.optim.AdamW(optim_dict, betas=(0.9, 0.95))
        if opt.optim.sched: 
            self.sched = torch.optim.lr_scheduler.CosineAnnealingLR(self.optim, opt.max_epoch)
        if opt.optim.amp:
            self.scaler = torch.cuda.amp.GradScaler()

    def restore_checkpoint(self, opt, best=False, evaluate=False):
        epoch_start, iter_start = None, None
        if opt.resume:
            if opt.device == 0: print("resuming from previous checkpoint...")
            epoch_start, iter_start, best_val, best_ep = util.restore_checkpoint(opt, self, resume=opt.resume, best=best, evaluate=evaluate)
            self.best_val = best_val
            self.best_ep = best_ep
        elif opt.load is not None:
            if opt.device == 0: print("loading weights from checkpoint {}...".format(opt.load))
            epoch_start, iter_start, best_val, best_ep = util.restore_checkpoint(opt, self, load_name=opt.load)
        else:
            if opt.device == 0: print("initializing weights from scratch...")
        self.epoch_start = epoch_start or 0
        self.iter_start = iter_start or 0

    def setup_visualizer(self, opt, test=False):
        if opt.device == 0: 
            print("setting up visualizers...")
            if opt.tb:
                self.tb = torch.utils.tensorboard.SummaryWriter(log_dir=opt.output_path, flush_secs=10)
    
    def train(self, opt):
        # before training
        torch.cuda.set_device(opt.device)
        torch.cuda.empty_cache()
        if opt.device == 0: print("TRAINING START")
        self.train_metric_logger = util.MetricLogger(delimiter="  ")
        self.train_metric_logger.add_meter('lr', util.SmoothedValue(window_size=1, fmt='{value:.6f}'))
        self.iter_skip = self.iter_start % len(self.train_loader)
        self.it = self.iter_start
        self.skip_dis = False
        if not opt.resume: 
            self.best_val = np.inf
            self.best_ep = 1
        # training
        if self.iter_start == 0 and not opt.debug: self.evaluate(opt, ep=0, training=True)
        # if opt.device == 0: self.save_checkpoint(opt, ep=0, it=0, best_val=self.best_val, best_ep=self.best_ep)
        self.ep = self.epoch_start
        for self.ep in range(self.epoch_start, opt.max_epoch):
            self.train_epoch(opt)
        # after training
        if opt.device == 0: self.save_checkpoint(opt, ep=self.ep, it=self.it, best_val=self.best_val, best_ep=self.best_ep)
        if opt.tb and opt.device == 0:
            self.tb.flush()
            self.tb.close()
        if opt.device == 0: 
            print("TRAINING DONE")
            print("Best val: %.4f @ epoch %d" % (self.best_val, self.best_ep))
        cleanup()
    
    def train_epoch(self, opt):
        # before train epoch
        self.train_loader.sampler.set_epoch(self.ep)
        if opt.device == 0:
            print("training epoch {}".format(self.ep+1))
        batch_progress = range(self.num_batches)
        self.graph.train()
        # train epoch
        loader = iter(self.train_loader)
           
        for batch_id in batch_progress:
            # if resuming from previous checkpoint, skip until the last iteration number is reached
            if self.iter_skip>0:
                self.iter_skip -= 1
                continue
            batch = next(loader)
            # train iteration
            var = edict(batch)
            opt.H, opt.W = opt.image_size
            var = util.move_to_device(var, opt.device)
            loss = self.train_iteration(opt, var, batch_progress)
            
        # after train epoch
        lr = self.sched.get_last_lr()[0] if opt.optim.sched else opt.optim.lr
        if opt.optim.sched: self.sched.step()
        if (self.ep + 1) % opt.freq.eval == 0:
            if opt.device == 0: print("validating epoch {}".format(self.ep+1))
            current_val = self.evaluate(opt, ep=self.ep+1, training=True)
            if current_val < self.best_val and opt.device == 0:
                self.best_val = current_val
                self.best_ep = self.ep + 1
                self.save_checkpoint(opt, ep=self.ep, it=self.it, best_val=self.best_val, best_ep=self.best_ep, best=True, latest=True)

    def train_iteration(self, opt, var, loader):
        # before train iteration
        torch.distributed.barrier()
        
        # train iteration
        with torch.autocast(device_type='cuda', dtype=torch.float16, enabled=opt.optim.amp):
            var, loss = self.graph.forward(opt, var, training=True, get_loss=True)
            loss = self.summarize_loss(opt, var, loss)
            loss_scaled = loss.all / opt.optim.accum
            
        # backward
        if opt.optim.amp:
            self.scaler.scale(loss_scaled).backward()
            # skip update if accumulating gradient
            if (self.it + 1) % opt.optim.accum == 0:
                self.scaler.unscale_(self.optim)
                # gradient clipping
                if opt.optim.clip_norm:
                    norm = torch.nn.utils.clip_grad_norm_(self.graph.parameters(), opt.optim.clip_norm)
                    if opt.debug: print("Grad norm: {}".format(norm))
                self.scaler.step(self.optim)
                self.scaler.update()
                self.optim.zero_grad()
        else:
            loss_scaled.backward()
            if (self.it + 1) % opt.optim.accum == 0:
                if opt.optim.clip_norm:
                    norm = torch.nn.utils.clip_grad_norm_(self.graph.parameters(), opt.optim.clip_norm)
                    if opt.debug: print("Grad norm: {}".format(norm))
                self.optim.step()
                self.optim.zero_grad()

        # after train iteration
        lr = self.sched.get_last_lr()[0] if opt.optim.sched else opt.optim.lr
        self.train_metric_logger.update(lr=lr)
        self.train_metric_logger.update(loss=loss.all)
        if opt.device == 0: 
            if (self.it) % opt.freq.vis == 0 and not opt.debug: 
                self.visualize(opt, var, step=self.it, split="train")
            if (self.it+1) % opt.freq.ckpt_latest == 0 and not opt.debug: 
                self.save_checkpoint(opt, ep=self.ep, it=self.it+1, best_val=self.best_val, best_ep=self.best_ep, latest=True)
            if (self.it) % opt.freq.scalar == 0 and not opt.debug: 
                self.log_scalars(opt, var, loss, step=self.it, split="train")
            if (self.it) % (opt.freq.save_vis * (self.it//10000*10+1)) == 0 and not opt.debug: 
                self.vis_train_iter(opt)
            if (self.it) % opt.freq.print == 0:
                print('[{}] '.format(datetime.datetime.now().time()), end='')
                print(f'Train Iter {self.it}/{self.num_batches*opt.max_epoch}: {self.train_metric_logger}')
        self.it += 1
        return loss
        
    @torch.no_grad()
    def vis_train_iter(self, opt):
        self.graph.eval()
        for i in range(len(self.viz_data)):
            var_viz = edict(deepcopy(self.viz_data[i]))
            var_viz = util.move_to_device(var_viz, opt.device)
            var_viz = self.graph.module(opt, var_viz, training=False, get_loss=False)
            vis_folder = "vis_log/iter_{}".format(self.it)
            os.makedirs("{}/{}".format(opt.output_path, vis_folder), exist_ok=True)
            util_vis.dump_images(opt, var_viz.idx, "image_input", var_viz.rgb_input_map, masks=None, from_range=(0, 1), folder=vis_folder)
            util_vis.dump_images(opt, var_viz.idx, "mask_input", var_viz.mask_input_map, folder=vis_folder)
            util_vis.dump_depths(opt, var_viz.idx, "depth_est", var_viz.depth_pred, var_viz.mask_input_map, rescale=True, folder=vis_folder)
            util_vis.dump_depths(opt, var_viz.idx, "depth_input", var_viz.depth_input_map, var_viz.mask_input_map, rescale=True, folder=vis_folder)
            if 'seen_points_pred' in var_viz and 'seen_points_gt' in var_viz:
                util_vis.dump_pointclouds_compare(opt, var_viz.idx, "seen_surface", var_viz.seen_points_pred, var_viz.seen_points_gt, folder=vis_folder)
        self.graph.train()

    def summarize_loss(self, opt, var, loss, non_act_loss_key=[]):
        loss_all = 0.
        assert("all" not in loss)
        # weigh losses
        for key in loss:
            assert(key in opt.loss_weight)
            if opt.loss_weight[key] is not None:
                assert not torch.isinf(loss[key].mean()), "loss {} is Inf".format(key)
                assert not torch.isnan(loss[key].mean()), "loss {} is NaN".format(key)
                loss_all += float(opt.loss_weight[key])*loss[key].mean() if key not in non_act_loss_key else 0.0*loss[key].mean()
        loss.update(all=loss_all)
        return loss

# =================================================== set up evaluation =========================================================

    @torch.no_grad()
    def evaluate(self, opt, ep, training=False):
        self.graph.eval()
        loss_eval = edict()
        
        # metric dictionary
        metric_eval = {}
        for metric_key in self.depth_metric.metric_keys:
            metric_eval[metric_key] = []
        metric_avg = {}
        eval_metric_logger = util.MetricLogger(delimiter="  ")
        
        # dataloader on the test set
        with torch.cuda.device(opt.device):
            for it, batch in enumerate(self.test_loader):
                
                # inference the model
                var = edict(batch)
                var = self.evaluate_batch(opt, var, ep, it, single_gpu=False)

                # record foreground mae for evaluation
                sample_metrics, var.depth_pred_aligned = self.depth_metric.compute_metrics(
                    var.depth_pred, var.depth_input_map, var.mask_eroded if 'mask_eroded' in var else var.mask_input_map)
                var.rmse = sample_metrics['rmse']
                curr_metrics = {}
                for metric_key in metric_eval:
                    metric_eval[metric_key].append(sample_metrics[metric_key])
                    curr_metrics[metric_key] = sample_metrics[metric_key].mean()
                eval_metric_logger.update(**curr_metrics)
                # eval_metric_logger.update(metric_key=sample_metrics[metric_key].mean())
                
                # accumulate the scores
                if opt.device == 0 and it % opt.freq.print_eval == 0: 
                    print('[{}] '.format(datetime.datetime.now().time()), end='')
                    print(f'Eval Iter {it}/{len(self.test_loader)} @ EP {ep}: {eval_metric_logger}')
                
                # dump the result if in eval mode
                if not training: 
                    self.dump_results(opt, var, ep, write_new=(it == 0))

                # save the visualization
                if it == 0 and training and opt.device == 0: 
                    print("visualizing and saving results...")
                    for i in range(len(self.viz_data)):
                        var_viz = edict(deepcopy(self.viz_data[i]))
                        var_viz = self.evaluate_batch(opt, var_viz, ep, it, single_gpu=True)
                        self.visualize(opt, var_viz, step=ep, split="eval")
                        self.dump_results(opt, var_viz, ep, train=True)
            
            # collect the eval results into tensors
            for metric_key in metric_eval:
                metric_eval[metric_key] = torch.cat(metric_eval[metric_key], dim=0)

        if opt.world_size > 1:
            metric_gather_dict = {}
            # empty tensors for gathering
            for metric_key in metric_eval:
                metric_gather_dict[metric_key] = [torch.zeros_like(metric_eval[metric_key]).to(opt.device) for _ in range(opt.world_size)]
            
            # gather the metrics
            torch.distributed.barrier()
            for metric_key in metric_eval:
                torch.distributed.all_gather(metric_gather_dict[metric_key], metric_eval[metric_key])
                metric_gather_dict[metric_key] = torch.cat(metric_gather_dict[metric_key], dim=0)
        else:
            metric_gather_dict = metric_eval
            
        # handle last batch, if any
        if len(self.test_loader.sampler) * opt.world_size < len(self.test_data):
            for metric_key in metric_eval:
                metric_gather_dict[metric_key] = [metric_gather_dict[metric_key]]
            for batch in self.aux_test_loader:
                # inference the model
                var = edict(batch)
                var = self.evaluate_batch(opt, var, ep, it, single_gpu=False)

                # record MAE for evaluation
                sample_metrics, var.depth_pred_aligned = self.depth_metric.compute_metrics(
                    var.depth_pred, var.depth_input_map, var.mask_eroded if 'mask_eroded' in var else var.mask_input_map)
                var.rmse = sample_metrics['rmse']
                for metric_key in metric_eval:
                    metric_gather_dict[metric_key].append(sample_metrics[metric_key])
                
                # dump the result if in eval mode
                if not training and opt.device == 0: 
                    self.dump_results(opt, var, ep, write_new=(it == 0))
            
            for metric_key in metric_eval:
                metric_gather_dict[metric_key] = torch.cat(metric_gather_dict[metric_key], dim=0)

        assert metric_gather_dict['l1_err'].shape[0] == len(self.test_data)
        # compute the mean of the metrics
        for metric_key in metric_eval:
            metric_avg[metric_key] = metric_gather_dict[metric_key].mean()
        
        # printout and save the metrics     
        if opt.device == 0:
            # print eval info
            print_eval(opt, depth_metrics=metric_avg)
            val_metric = metric_avg['l1_err']
        
            if training:
                # log/visualize results to tb/vis
                self.log_scalars(opt, var, loss_eval, metric=metric_avg, step=ep, split="eval")
            
            if not training:
                # write to file
                metrics_file = os.path.join(opt.output_path, 'best_val.txt')
                with open(metrics_file, "w") as outfile:
                    for metric_key in metric_avg:
                        outfile.write('{}: {:.6f}\n'.format(metric_key, metric_avg[metric_key].item()))
        
            return val_metric.item()
        return float('inf')

    def evaluate_batch(self, opt, var, ep=None, it=None, single_gpu=False):
        var = util.move_to_device(var, opt.device)
        if single_gpu:
            var  = self.graph.module(opt, var, training=False, get_loss=False)
        else:
            var = self.graph(opt, var, training=False, get_loss=False)
        return var

    @torch.no_grad()
    def log_scalars(self, opt, var, loss, metric=None, step=0, split="train"):
        if split=="train":
            sample_metrics, _ = self.depth_metric.compute_metrics(
                var.depth_pred, var.depth_input_map, var.mask_eroded if 'mask_eroded' in var else var.mask_input_map)
            metric = dict(L1_ERR=sample_metrics['l1_err'].mean().item())
        for key, value in loss.items():
            if key=="all": continue
            self.tb.add_scalar("{0}/loss_{1}".format(split, key), value.mean(), step)
        if metric is not None:
            for key, value in metric.items():
                self.tb.add_scalar("{0}/{1}".format(split, key), value, step)

    @torch.no_grad()
    def visualize(self, opt, var, step=0, split="train"):
        pass

    @torch.no_grad()
    def dump_results(self, opt, var, ep, write_new=False, train=False):
        # create the dir
        current_folder = "dump" if train == False else "vis_{}".format(ep)
        os.makedirs("{}/{}/".format(opt.output_path, current_folder), exist_ok=True)
        
        # save the results
        util_vis.dump_images(opt, var.idx, "image_input", var.rgb_input_map, masks=None, from_range=(0, 1), folder=current_folder)
        util_vis.dump_images(opt, var.idx, "mask_input", var.mask_input_map, folder=current_folder)
        util_vis.dump_depths(opt, var.idx, "depth_pred", var.depth_pred, var.mask_input_map, rescale=True, folder=current_folder)
        util_vis.dump_depths(opt, var.idx, "depth_input", var.depth_input_map, var.mask_input_map, rescale=True, folder=current_folder)
        if 'seen_points_pred' in var and 'seen_points_gt' in var:
            util_vis.dump_pointclouds_compare(opt, var.idx, "seen_surface", var.seen_points_pred, var.seen_points_gt, folder=current_folder)
        
        if "depth_pred_aligned" in var:
            # get the max and min for the depth map
            batch_size = var.depth_input_map.shape[0]
            mask = var.mask_eroded if 'mask_eroded' in var else var.mask_input_map
            masked_depth_far_bg = var.depth_input_map * mask + (1 - mask) * 1000
            depth_min_gt = masked_depth_far_bg.view(batch_size, -1).min(dim=1)[0]
            masked_depth_invalid_bg = var.depth_input_map * mask + (1 - mask) * 0
            depth_max_gt = masked_depth_invalid_bg.view(batch_size, -1).max(dim=1)[0]
            depth_vis_pred = (var.depth_pred_aligned - depth_min_gt.view(batch_size, 1, 1, 1)) / (depth_max_gt - depth_min_gt).view(batch_size, 1, 1, 1)
            depth_vis_pred = depth_vis_pred * mask + (1 - mask)
            depth_vis_gt = (var.depth_input_map - depth_min_gt.view(batch_size, 1, 1, 1)) / (depth_max_gt - depth_min_gt).view(batch_size, 1, 1, 1)
            depth_vis_gt = depth_vis_gt * mask + (1 - mask)
            util_vis.dump_depths(opt, var.idx, "depth_gt_aligned", depth_vis_gt.clamp(max=1, min=0), None, rescale=False, folder=current_folder)
            util_vis.dump_depths(opt, var.idx, "depth_pred_aligned", depth_vis_pred.clamp(max=1, min=0), None, rescale=False, folder=current_folder)
            if "mask_eroded" in var and "rmse" in var:
                util_vis.dump_images(opt, var.idx, "image_eroded", var.rgb_input_map, masks=var.mask_eroded, metrics=var.rmse, from_range=(0, 1), folder=current_folder)

    def save_checkpoint(self, opt, ep=0, it=0, best_val=np.inf, best_ep=1, latest=False, best=False):
        util.save_checkpoint(opt, self, ep=ep, it=it, best_val=best_val, best_ep=best_ep, latest=latest, best=best)
        if not latest:
            print("checkpoint saved: ({0}) {1}, epoch {2} (iteration {3})".format(opt.group, opt.name, ep, it))
        if best:
            print("Saving the current model as the best...")