Spaces:
Sleeping
Sleeping
| import csv | |
| import datetime | |
| from collections import defaultdict | |
| import numpy as np | |
| import torch | |
| import torchvision | |
| import wandb | |
| from termcolor import colored | |
| from torch.utils.tensorboard import SummaryWriter | |
| COMMON_TRAIN_FORMAT = [('frame', 'F', 'int'), ('step', 'S', 'int'), | |
| ('episode', 'E', 'int'), ('episode_length', 'L', 'int'), | |
| ('episode_reward', 'R', 'float'), | |
| ('fps', 'FPS', 'float'), ('total_time', 'T', 'time')] | |
| COMMON_EVAL_FORMAT = [('frame', 'F', 'int'), ('step', 'S', 'int'), | |
| ('episode', 'E', 'int'), ('episode_length', 'L', 'int'), | |
| ('episode_reward', 'R', 'float'), | |
| ('total_time', 'T', 'time')] | |
| class AverageMeter(object): | |
| def __init__(self): | |
| self._sum = 0 | |
| self._count = 0 | |
| def update(self, value, n=1): | |
| self._sum += value | |
| self._count += n | |
| def value(self): | |
| return self._sum / max(1, self._count) | |
| class MetersGroup(object): | |
| def __init__(self, csv_file_name, formating, use_wandb): | |
| self._csv_file_name = csv_file_name | |
| self._formating = formating | |
| self._meters = defaultdict(AverageMeter) | |
| self._csv_file = None | |
| self._csv_writer = None | |
| self.use_wandb = use_wandb | |
| def log(self, key, value, n=1): | |
| self._meters[key].update(value, n) | |
| def _prime_meters(self): | |
| data = dict() | |
| for key, meter in self._meters.items(): | |
| if key.startswith('train'): | |
| key = key[len('train') + 1:] | |
| else: | |
| key = key[len('eval') + 1:] | |
| key = key.replace('/', '_') | |
| data[key] = meter.value() | |
| return data | |
| def _remove_old_entries(self, data): | |
| rows = [] | |
| with self._csv_file_name.open('r') as f: | |
| reader = csv.DictReader(f) | |
| for row in reader: | |
| if 'episode' in row: | |
| # BUGFIX: covers weird cases where CSV are badly written | |
| if row['episode'] == '': | |
| rows.append(row) | |
| continue | |
| if type(row['episode']) == type(None): | |
| continue | |
| if float(row['episode']) >= data['episode']: | |
| break | |
| rows.append(row) | |
| with self._csv_file_name.open('w') as f: | |
| # To handle CSV that have more keys than new data | |
| keys = set(data.keys()) | |
| if len(rows) > 0: keys = keys | set(row.keys()) | |
| keys = sorted(list(keys)) | |
| # | |
| writer = csv.DictWriter(f, | |
| fieldnames=keys, | |
| restval=0.0) | |
| writer.writeheader() | |
| for row in rows: | |
| writer.writerow(row) | |
| def _dump_to_csv(self, data): | |
| if self._csv_writer is None: | |
| should_write_header = True | |
| if self._csv_file_name.exists(): | |
| self._remove_old_entries(data) | |
| should_write_header = False | |
| self._csv_file = self._csv_file_name.open('a') | |
| self._csv_writer = csv.DictWriter(self._csv_file, | |
| fieldnames=sorted(data.keys()), | |
| restval=0.0) | |
| if should_write_header: | |
| self._csv_writer.writeheader() | |
| # To handle components that start training later | |
| # (restval covers only when data has less keys than the CSV) | |
| if self._csv_writer.fieldnames != sorted(data.keys()) and \ | |
| len(self._csv_writer.fieldnames) < len(data.keys()): | |
| self._csv_file.close() | |
| self._csv_file = self._csv_file_name.open('r') | |
| dict_reader = csv.DictReader(self._csv_file) | |
| rows = [row for row in dict_reader] | |
| self._csv_file.close() | |
| self._csv_file = self._csv_file_name.open('w') | |
| self._csv_writer = csv.DictWriter(self._csv_file, | |
| fieldnames=sorted(data.keys()), | |
| restval=0.0) | |
| self._csv_writer.writeheader() | |
| for row in rows: | |
| self._csv_writer.writerow(row) | |
| self._csv_writer.writerow(data) | |
| self._csv_file.flush() | |
| def _format(self, key, value, ty): | |
| if ty == 'int': | |
| value = int(value) | |
| return f'{key}: {value}' | |
| elif ty == 'float': | |
| return f'{key}: {value:.04f}' | |
| elif ty == 'time': | |
| value = str(datetime.timedelta(seconds=int(value))) | |
| return f'{key}: {value}' | |
| else: | |
| raise f'invalid format type: {ty}' | |
| def _dump_to_console(self, data, prefix): | |
| prefix = colored(prefix, 'yellow' if prefix == 'train' else 'green') | |
| pieces = [f'| {prefix: <14}'] | |
| for key, disp_key, ty in self._formating: | |
| value = data.get(key, 0) | |
| pieces.append(self._format(disp_key, value, ty)) | |
| print(' | '.join(pieces)) | |
| def _dump_to_wandb(self, data): | |
| wandb.log(data) | |
| def dump(self, step, prefix): | |
| if len(self._meters) == 0: | |
| return | |
| data = self._prime_meters() | |
| data['frame'] = step | |
| if self.use_wandb: | |
| wandb_data = {prefix + '/' + key: val for key, val in data.items()} | |
| self._dump_to_wandb(data=wandb_data) | |
| # self._dump_to_csv(data) | |
| self._dump_to_console(data, prefix) | |
| self._meters.clear() | |
| class Logger(object): | |
| def __init__(self, log_dir, use_tb, use_wandb): | |
| self._log_dir = log_dir | |
| self._train_mg = MetersGroup(log_dir / 'train.csv', | |
| formating=COMMON_TRAIN_FORMAT, | |
| use_wandb=use_wandb) | |
| self._eval_mg = MetersGroup(log_dir / 'eval.csv', | |
| formating=COMMON_EVAL_FORMAT, | |
| use_wandb=use_wandb) | |
| if use_tb: | |
| self._sw = SummaryWriter(str(log_dir / 'tb')) | |
| else: | |
| self._sw = None | |
| self.use_wandb = use_wandb | |
| def _try_sw_log(self, key, value, step): | |
| if self._sw is not None: | |
| self._sw.add_scalar(key, value, step) | |
| def log(self, key, value, step): | |
| assert key.startswith('train') or key.startswith('eval') | |
| if type(value) == torch.Tensor: | |
| value = value.item() | |
| self._try_sw_log(key, value, step) | |
| mg = self._train_mg if key.startswith('train') else self._eval_mg | |
| mg.log(key, value) | |
| def log_metrics(self, metrics, step, ty): | |
| for key, value in metrics.items(): | |
| self.log(f'{ty}/{key}', value, step) | |
| def dump(self, step, ty=None): | |
| if ty is None or ty == 'eval': | |
| self._eval_mg.dump(step, 'eval') | |
| if ty is None or ty == 'train': | |
| self._train_mg.dump(step, 'train') | |
| def log_and_dump_ctx(self, step, ty): | |
| return LogAndDumpCtx(self, step, ty) | |
| def log_visual(self, data, step): | |
| if self._sw is not None: | |
| for k, v in data.items(): | |
| if len(v.shape) == 3: | |
| self._sw.add_image(k, v) | |
| else: | |
| if len(v.shape) == 4: | |
| v = np.expand_dims(v, axis=0) | |
| self._sw.add_video(k, v, global_step=step, fps=15) | |
| if self.use_wandb: | |
| for k, v in data.items(): | |
| if type(v) is not np.ndarray: | |
| v = v.cpu() | |
| if v.dtype not in [np.uint8]: | |
| v = v*255 | |
| v = np.uint8(v) | |
| if len(v.shape) == 3: | |
| if v.shape[0] == 3: | |
| v = v.transpose(1,2,0) | |
| # Note: defaulting to save only one image/video to save storage on wandb | |
| wandb.log({k: wandb.Image(v)},) | |
| else: | |
| # Note: defaulting to save only one image/video to save storage on wandb | |
| wandb.log({k: wandb.Video(v, fps=15, format="gif")},) | |
| class LogAndDumpCtx: | |
| def __init__(self, logger, step, ty): | |
| self._logger = logger | |
| self._step = step | |
| self._ty = ty | |
| def __enter__(self): | |
| return self | |
| def __call__(self, key, value): | |
| self._logger.log(f'{self._ty}/{key}', value, self._step) | |
| def __exit__(self, *args): | |
| self._logger.dump(self._step, self._ty) | |