# Copyright 2024 Toyota Research Institute. All rights reserved. import torch import torch.distributed as dist import matplotlib.pyplot as plt import numpy as np import wandb from tqdm import tqdm def save_close(logs, fig, name): logs[name] = wandb.Image(fig) # fig.savefig(name) plt.close(fig) class Evaluator: """ Wandb logger class to monitor training. Parameters ---------- """ def __init__(self, cfg): from custom.eval.depth import DepthEvaluation from custom.eval.rgb import RGBEvaluation # NOTE(bvh): use vidar's Config here (not custom/utils/config.py:Config) from vidar.utils.config import Config self.metrics = {} for key, val in cfg.items(): if key.startswith('depth'): self.metrics[key] = DepthEvaluation(Config(**val)) if key.startswith('rgb'): self.metrics[key] = RGBEvaluation(Config(**val)) self.length = None self.results = None self.detail_level = cfg.get('detail_level', 4) self.gt_source = cfg.get('gt_source', 'gt') self.modes = cfg.get('modes', ['r','rpv','a','apv','i','ipv']) # 'mpc','mpt','mptc','a','atpc','acpt','itc','icat','itac','i'] def draw_plot(self, name, gathered, i, xlabel, mode): # name = self.metrics[key].metrics[i] # Metric name for plot if mode == 'mean': x = [str(k) for k in gathered.keys()] # Use keys as the X labels y = [v[i].detach().cpu().numpy() for v in gathered.values()] # Use metric values as Y labels elif mode == 'indiv': x = gathered[:, 0].detach().cpu().numpy() # Use keys as the X labels y = gathered[:, i+1].detach().cpu().numpy() # Use metric values as Y labels else: raise ValueError('Invalid mode') y_avg = [np.mean(y)] * len(x) # Get average for dotted line fig, ax = plt.subplots(figsize =(32, 8)) plt.plot(x, y, 'k-D', linewidth=2, markersize=3) plt.plot(x, y_avg, 'r--', linewidth=3) plt.xlabel(xlabel, fontsize=18) plt.ylabel(name, fontsize=18) if mode == 'indiv': plt.xticks(np.arange(min(x), max(x) + 1, 10.0)) plt.grid() plt.xticks(fontsize=14) plt.yticks(fontsize=14) # Annotate dots with values for (xx, yy) in zip(x, y): ax.annotate(f'%.03f' % yy, xy=(xx,yy), fontsize=15, xytext=(2,0), textcoords='offset points', color='b') ax.annotate(f'%.03f' % y_avg[0], xy=(x[0],y_avg[0]), fontsize=16, xytext=(-40,0), textcoords='offset points', color='r') # Tight borders fig.tight_layout() return fig def process_gathered(self, wandb_logs, all_gathered, all_gathered_mean, task, dataloader_key=None): prefix = 'metrics/val' if dataloader_key is not None: prefix = f'{prefix}/{dataloader_key}' prefix = f'{prefix}/{task}' keys = all_gathered_mean.keys() time = sorted(list(set(key[0] for key in keys))) view = sorted(list(set(key[1] for key in keys))) # indiv_per_time_view = all_gathered # indiv_per_view = {kc: {kt[0]: vt for kt, vt in indiv_per_time_view.items() if kt[1] == kc} for kc in view} # indiv_per_time = {kt: {kc[1]: vc for kc, vc in indiv_per_time_view.items() if kc[0] == kt} for kt in time} # indiv_reduced_per_time_view = torch.stack([v for v in indiv_per_time_view.values()], 0).mean(0) # indiv_reduced_per_view = {kc: torch.stack([v for v in vc.values()], 0).mean(0) for kc, vc in indiv_per_view.items()} # indiv_reduced_per_time = {kt: torch.stack([v for v in vt.values()], 0).mean(0) for kt, vt in indiv_per_time.items()} per_time_view = all_gathered_mean per_view = {kc: {kt[0]: vt for kt, vt in per_time_view.items() if kt[1] == kc} for kc in view} per_time = {kt: {kc[1]: vc for kc, vc in per_time_view.items() if kc[0] == kt} for kt in time} reduced = torch.stack([v for v in per_time_view.values()], 0).mean(0) reduced_per_view = {kc: torch.stack([v for v in vc.values()], 0).mean(0) for kc, vc in per_view.items()} reduced_per_time = {kt: torch.stack([v for v in vt.values()], 0).mean(0) for kt, vt in per_time.items()} # Sample counts per (time, cam): each row in all_gathered[tc] is one observed # sample index, so .shape[0] = # samples that contributed to that (t, c). n_per_time_view = {tc: int(g.shape[0]) for tc, g in all_gathered.items()} n_per_view = {kc: sum(n for tc, n in n_per_time_view.items() if tc[1] == kc) for kc in view} n_per_time = {kt: sum(n for tc, n in n_per_time_view.items() if tc[0] == kt) for kt in time} n_reduced = sum(n_per_time_view.values()) metrics = self.metrics[task.split('/')[0]].metrics for i, metric in enumerate(metrics): if 'reduced' in self.modes or 'r' in self.modes: wandb_logs[f'{prefix}/reduced/{metric}'] = reduced[i] if i == 0: wandb_logs[f'{prefix}/reduced/count'] = n_reduced if 'reduced_per_view' in self.modes or 'rpv' in self.modes: for key in tqdm(reduced_per_view.keys(), desc='reduced_per_view', leave=False, ncols=96): wandb_logs[f'{prefix}/reduced/per_view/{key}/{metric}'] = reduced_per_view[key][i] if i == 0: wandb_logs[f'{prefix}/reduced/per_view/{key}/count'] = n_per_view[key] if 'reduced_per_time' in self.modes or 'rpt' in self.modes: for key in tqdm(reduced_per_time.keys(), desc='reduced_per_time', leave=False, ncols=96): wandb_logs[f'{prefix}/reduced/per_time/{key}/{metric}'] = reduced_per_time[key][i] if i == 0: wandb_logs[f'{prefix}/reduced/per_time/{key}/count'] = n_per_time[key] if 'reduced_per_timeview' in self.modes or 'rptv' in self.modes: for key in tqdm(per_time_view.keys(), desc='reduced_per_timeview', leave=False, ncols=96): wandb_logs[f'{prefix}/reduced/per_timeview/{key[0]}_{key[1]}/{metric}'] = per_time_view[key][i] if i == 0: wandb_logs[f'{prefix}/reduced/per_timeview/{key[0]}_{key[1]}/count'] = n_per_time_view[key] if 'avg' in self.modes or 'a' in self.modes: fig = self.draw_plot(metric, per_time_view, i, xlabel=f'Average (Timestep and View)', mode='mean') save_close(wandb_logs, fig, f'{prefix}/average/{metric}') if 'avg_per_view' in self.modes or 'apv' in self.modes: for key in tqdm(per_view.keys(), desc='avg_per_view', leave=False, ncols=96): fig = self.draw_plot(metric, per_view[key], i, xlabel=f'Average - Timestep (View {key})', mode='mean') save_close(wandb_logs, fig, f'{prefix}/average/per_view/{key}/{metric}') if 'avg_per_time' in self.modes or 'apt' in self.modes: for key in tqdm(per_time.keys(), desc='avg_per_time', leave=False, ncols=96): fig = self.draw_plot(metric, per_time[key], i, xlabel=f'Average - View (Timestep {key})', mode='mean') save_close(wandb_logs, fig, f'{prefix}/average/per_time/{key}/{metric}') # if 'indiv_per_view' in self.modes or 'ipv' in self.modes: # for key in tqdm(indiv_reduced_per_view.keys(), desc='indiv_per_view', leave=False, ncols=96): # fig = self.draw_plot(metric, indiv_reduced_per_view[key], i, # xlabel=f'Individual - Average Timestep (View {key})', mode='indiv') # save_close(wandb_logs, fig, f'{prefix}/individual/per_view/{key}/{metric}') # if 'indiv_per_time' in self.modes or 'ipt' in self.modes: # for key in tqdm(indiv_reduced_per_time.keys(), desc='indiv_per_time', leave=False, ncols=96): # fig = self.draw_plot(metric, indiv_reduced_per_time[key], i, # xlabel=f'Individual - Average View (Timestep {key})', mode='indiv') # save_close(wandb_logs, fig, f'{prefix}/individual/per_time/{key}/{metric}') # if 'indiv_per_timeview' in self.modes or 'iptv' in self.modes: # for key in tqdm(indiv_per_time_view.keys(), desc='indiv_per_timeview', leave=False, ncols=96): # fig = self.draw_plot(metric, indiv_per_time_view[key], i, # xlabel=f'Individual (Timestep {key[0]} View {key[1]})', mode='indiv') # save_close(wandb_logs, fig, f'{prefix}/individual/per_timeview/{key[0]}_{key[1]}/{metric}') # if 'indiv' in self.modes or 'i' in self.modes: # fig = self.draw_plot(metric, indiv_reduced_per_time_view, i, # xlabel=f'Individual - Average Timestep and View', mode='indiv') # save_close(wandb_logs, fig, f'{prefix}/individual/{metric}') return wandb_logs def prepare(self, dataloader): self.length = len(dataloader) * dist.get_world_size() self.results = {} def step(self, batch, output): # Get index and dataset name for logging idx, dset = batch['dl_idx'], batch['dset_name'] # Retrieve unrolled dicts from model (may not exist for configs with low-dim streams) if 'unrolled' not in output: return # Skip evaluation when unrolled is not available unrolled = output['unrolled'] gtruths, predictions = unrolled['gt'], unrolled['pred'] # Compare to reconstructions if requested if self.gt_source == 'rec': gtruths = unrolled['rec'] # Loop over all metrics for key, val in self.metrics.items(): # Loop over all predictions for predictions_key, predictions_val in predictions.items(): # If prediction starts with the metric name if predictions_key.startswith(key.split('_')[0]): metrics_predictions_key = f'{key}/{predictions_key}' # Loop over all timesteps and cameras for time_cam in predictions_val.keys(): gt = gtruths[predictions_key][time_cam] # Ground-truth pred = predictions[predictions_key][time_cam] # Prediction calculated = val.evaluate(gt, pred) # Metric # Create dictionary if hasn't been done before if not metrics_predictions_key in self.results.keys(): self.results[metrics_predictions_key] = {} # Create results matrix if it's not there already key_time_cam = f'{metrics_predictions_key}-{str(time_cam).replace(" ", "")}' if not key_time_cam in self.results[metrics_predictions_key].keys(): self.results[metrics_predictions_key][key_time_cam] = torch.zeros( (self.length, len(val.metrics) + 2), device='cuda', dtype=torch.float64) # In case batch size > 1 for i in range(len(idx)): if calculated[i].mean() != -999: # Concatenate [1, idx] to metrics for bookkeeping self.results[metrics_predictions_key][key_time_cam][idx[i]] = torch.cat( [torch.tensor([[1.0, idx[i]]], device=gt.device), calculated[[i]]], 1) else: calculated[i][:] = 0.0 self.results[metrics_predictions_key][key_time_cam][idx[i]] = torch.cat( [torch.tensor([[0.0, idx[i]]], device=gt.device), calculated[[i]]], 1) return None def finish(self, wandb_logs, dataloader_key): # Loop over all metrics for results_key1, results_val1 in self.results.items(): all_gathered = {} all_gathered_mean = {} # Loop over all timesteps and cameras for that metric for results_key, results_val in results_val1.items(): _, time_cam = results_key.split('-') time_cam = tuple([int(tc) for tc in time_cam[1:-1].split(',')]) # Create per-rank tensors for the gather operation and populate them gathered = [torch.zeros_like(results_val) for _ in range(dist.get_world_size())] dist.all_gather(gathered, results_val) gathered = sum(gathered) # Sum all ranks gathered = gathered[gathered[:, 0] > 0] # Remove indices not observed gathered = gathered[:, 1:] / gathered[:, [0]] # Divide by the number of times each index was observed gathered_mean = gathered[:, 1:].mean(0) # Get average of each metric # Calculate mean for each timestep and camera all_gathered[time_cam] = gathered all_gathered_mean[time_cam] = gathered_mean wandb_logs = self.process_gathered( wandb_logs, all_gathered, all_gathered_mean, task=results_key1, dataloader_key=dataloader_key ) return wandb_logs