ZeroShape / model /depth_engine.py
zxhuang1698's picture
initial commit
414b431
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...")