from typing import List, Optional import pathlib import pandas as pd import numpy as np import numba import click import time import collections import json import wandb import yaml import numbers import scipy.ndimage as sn from diffusion_policy.common.json_logger import read_json_log, JsonLogger import logging @numba.jit(nopython=True) def get_indexed_window_average( arr: np.ndarray, idxs: np.ndarray, window_size: int): result = np.zeros(idxs.shape, dtype=arr.dtype) length = arr.shape[0] for i in range(len(idxs)): idx = idxs[i] start = max(idx - window_size, 0) end = min(start + window_size, length) result[i] = np.mean(arr[start:end]) return result def compute_metrics(log_df: pd.DataFrame, key: str, end_step: Optional[int]=None, k_min_loss: int=10, k_around_max: int=10, max_k_window: int=10, replace_slash: int=True, ): if key not in log_df: return dict() # prepare data if end_step is not None: log_df = log_df.iloc[:end_step] is_key = ~pd.isnull(log_df[key]) is_key_idxs = is_key.index[is_key].to_numpy() if len(is_key_idxs) == 0: return dict() key_data = log_df[key][is_key].to_numpy() # after adding validation to workspace # rollout happens at the last step of each epoch # where the reported train_loss and val_loss # are already the average for that epoch train_loss = log_df['train_loss'][is_key].to_numpy() val_loss = log_df['val_loss'][is_key].to_numpy() result = dict() log_key = key if replace_slash: log_key = key.replace('/', '_') # max max_value = np.max(key_data) result['max/'+log_key] = max_value # k_around_max max_idx = np.argmax(key_data) end = min(max_idx + k_around_max // 2, len(key_data)) start = max(end - k_around_max, 0) k_around_max_value = np.mean(key_data[start:end]) result['k_around_max/'+log_key] = k_around_max_value # max_k_window k_window_value = sn.uniform_filter1d(key_data, size=max_k_window, axis=0, mode='nearest') max_k_window_value = np.max(k_window_value) result['max_k_window/'+log_key] = max_k_window_value # min_train_loss min_idx = np.argmin(train_loss) min_train_loss_value = key_data[min_idx] result['min_train_loss/'+log_key] = min_train_loss_value # min_val_loss min_idx = np.argmin(val_loss) min_val_loss_value = key_data[min_idx] result['min_val_loss/'+log_key] = min_val_loss_value # k_min_train_loss min_loss_idxs = np.argsort(train_loss)[:k_min_loss] k_min_train_loss_value = np.mean(key_data[min_loss_idxs]) result['k_min_train_loss/'+log_key] = k_min_train_loss_value # k_min_val_loss min_loss_idxs = np.argsort(val_loss)[:k_min_loss] k_min_val_loss_value = np.mean(key_data[min_loss_idxs]) result['k_min_val_loss/'+log_key] = k_min_val_loss_value # last result['last/'+log_key] = key_data[-1] # global step for visualization result['metric_global_step/'+log_key] = is_key_idxs[-1] return result def compute_metrics_agg( log_dfs: List[pd.DataFrame], key: str, end_step:int, **kwargs): # compute metrics results = collections.defaultdict(list) for log_df in log_dfs: result = compute_metrics(log_df, key=key, end_step=end_step, **kwargs) for k, v in result.items(): results[k].append(v) # agg agg_result = dict() for k, v in results.items(): value = np.mean(v) if k.startswith('metric_global_step'): value = int(value) agg_result[k] = value return agg_result @click.command() @click.option('--input', '-i', required=True, help='Root logging dir, contains train_* dirs') @click.option('--key', '-k', multiple=True, default=['test/mean_score']) @click.option('--interval', default=10, type=float) @click.option('--replace_slash', default=True, type=bool) @click.option('--index_key', '-ik', default='global_step') @click.option('--use_wandb', '-w', is_flag=True, default=False) @click.option('--project', default=None) @click.option('--name', default=None) @click.option('--id', default=None) @click.option('--group', default=None) def main( input, key, interval, replace_slash, index_key, use_wandb, # wandb args project, name, id, group): root_dir = pathlib.Path(input) assert root_dir.is_dir() metrics_dir = root_dir.joinpath('metrics') metrics_dir.mkdir(exist_ok=True) logging.basicConfig( level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s", handlers=[ logging.FileHandler(str(metrics_dir.joinpath("metrics.log"))), logging.StreamHandler() ] ) train_dirs = list(root_dir.glob('train_*')) log_files = [x.joinpath('logs.json.txt') for x in train_dirs] logging.info("Monitor waiting for log files!") while True: # wait for files to show up files_exist = True for log_file in log_files: if not log_file.is_file(): files_exist = False if files_exist: break time.sleep(1.0) logging.info("All log files ready!") # init path metric_log_path = metrics_dir.joinpath('logs.json.txt') metric_path = metrics_dir.joinpath('metrics.json') config_path = root_dir.joinpath('config.yaml') # load config config = yaml.safe_load(config_path.open('r')) # init wandb wandb_run = None if use_wandb: wandb_kwargs = config['logging'] if project is not None: wandb_kwargs['project'] = project if id is not None: wandb_kwargs['id'] = id if name is not None: wandb_kwargs['name'] = name if group is not None: wandb_kwargs['group'] = group wandb_kwargs['resume'] = True wandb_run = wandb.init( dir=str(metrics_dir), config=config, # auto-resume run, automatically load id # as long as using the same dir. # https://docs.wandb.ai/guides/track/advanced/resuming#resuming-guidance **wandb_kwargs ) wandb.config.update( { "output_dir": str(root_dir), } ) with JsonLogger(metric_log_path) as json_logger: last_log = json_logger.get_last_log() while True: # read json files log_dfs = [read_json_log(str(x), required_keys=key) for x in log_files] # previously logged data point last_log_idx = -1 if last_log is not None: last_log_idx = log_dfs[0].index[log_dfs[0][index_key] <= last_log[index_key]][-1] start_idx = last_log_idx + 1 # last idx where we have a data point from all logs end_idx = min(*[len(x) for x in log_dfs]) # log every position for this_idx in range(start_idx, end_idx): # compute metrics all_metrics = dict() global_step = log_dfs[0]['global_step'][this_idx] epoch = log_dfs[0]['epoch'][this_idx] all_metrics['global_step'] = global_step all_metrics['epoch'] = epoch for k in key: metrics = compute_metrics_agg( log_dfs=log_dfs, key=k, end_step=this_idx+1, replace_slash=replace_slash) all_metrics.update(metrics) # sanitize metrics old_metrics = all_metrics all_metrics = dict() for k, v in old_metrics.items(): if isinstance(v, numbers.Integral): all_metrics[k] = int(v) elif isinstance(v, numbers.Number): all_metrics[k] = float(v) has_update = all_metrics != last_log if has_update: last_log = all_metrics json_logger.log(all_metrics) with metric_path.open('w') as f: json.dump(all_metrics, f, sort_keys=True, indent=2) if wandb_run is not None: wandb_run.log(all_metrics, step=all_metrics[index_key]) logging.info(f"Metrics logged at step {all_metrics[index_key]}") time.sleep(interval) if __name__ == "__main__": main()