| import timeit |
| import numpy as np |
| import os |
| import os.path as osp |
| import shutil |
| import copy |
| import torch |
| import torch.nn as nn |
| import torch.distributed as dist |
| from .cfg_holder import cfg_unique_holder as cfguh |
| from . import sync |
|
|
| print_console_local_rank0_only = True |
|
|
| def print_log(*console_info): |
| local_rank = sync.get_rank('local') |
| if print_console_local_rank0_only and (local_rank!=0): |
| return |
| console_info = [str(i) for i in console_info] |
| console_info = ' '.join(console_info) |
| print(console_info) |
|
|
| if local_rank!=0: |
| return |
|
|
| log_file = None |
| try: |
| log_file = cfguh().cfg.train.log_file |
| except: |
| try: |
| log_file = cfguh().cfg.eval.log_file |
| except: |
| return |
| if log_file is not None: |
| with open(log_file, 'a') as f: |
| f.write(console_info + '\n') |
|
|
| class distributed_log_manager(object): |
| def __init__(self): |
| self.sum = {} |
| self.cnt = {} |
| self.time_check = timeit.default_timer() |
|
|
| cfgt = cfguh().cfg.train |
| use_tensorboard = getattr(cfgt, 'log_tensorboard', False) |
|
|
| self.ddp = sync.is_ddp() |
| self.rank = sync.get_rank('local') |
| self.world_size = sync.get_world_size('local') |
|
|
| self.tb = None |
| if use_tensorboard and (self.rank==0): |
| import tensorboardX |
| monitoring_dir = osp.join(cfguh().cfg.train.log_dir, 'tensorboard') |
| self.tb = tensorboardX.SummaryWriter(osp.join(monitoring_dir)) |
|
|
| def accumulate(self, n, **data): |
| if n < 0: |
| raise ValueError |
|
|
| for itemn, di in data.items(): |
| if itemn in self.sum: |
| self.sum[itemn] += di * n |
| self.cnt[itemn] += n |
| else: |
| self.sum[itemn] = di * n |
| self.cnt[itemn] = n |
|
|
| def get_mean_value_dict(self): |
| value_gather = [ |
| self.sum[itemn]/self.cnt[itemn] \ |
| for itemn in sorted(self.sum.keys()) ] |
|
|
| value_gather_tensor = torch.FloatTensor(value_gather).to(self.rank) |
| if self.ddp: |
| dist.all_reduce(value_gather_tensor, op=dist.ReduceOp.SUM) |
| value_gather_tensor /= self.world_size |
|
|
| mean = {} |
| for idx, itemn in enumerate(sorted(self.sum.keys())): |
| mean[itemn] = value_gather_tensor[idx].item() |
| return mean |
|
|
| def tensorboard_log(self, step, data, mode='train', **extra): |
| if self.tb is None: |
| return |
| if mode == 'train': |
| self.tb.add_scalar('other/epochn', extra['epochn'], step) |
| if 'lr' in extra: |
| self.tb.add_scalar('other/lr', extra['lr'], step) |
| for itemn, di in data.items(): |
| if itemn.find('loss') == 0: |
| self.tb.add_scalar('loss/'+itemn, di, step) |
| elif itemn == 'Loss': |
| self.tb.add_scalar('Loss', di, step) |
| else: |
| self.tb.add_scalar('other/'+itemn, di, step) |
| elif mode == 'eval': |
| if isinstance(data, dict): |
| for itemn, di in data.items(): |
| self.tb.add_scalar('eval/'+itemn, di, step) |
| else: |
| self.tb.add_scalar('eval', data, step) |
| return |
|
|
| def train_summary(self, itern, epochn, samplen, lr, tbstep=None): |
| console_info = [ |
| 'Iter:{}'.format(itern), |
| 'Epoch:{}'.format(epochn), |
| 'Sample:{}'.format(samplen),] |
|
|
| if lr is not None: |
| console_info += ['LR:{:.4E}'.format(lr)] |
|
|
| mean = self.get_mean_value_dict() |
|
|
| tbstep = itern if tbstep is None else tbstep |
| self.tensorboard_log( |
| tbstep, mean, mode='train', |
| itern=itern, epochn=epochn, lr=lr) |
|
|
| loss = mean.pop('Loss') |
| mean_info = ['Loss:{:.4f}'.format(loss)] + [ |
| '{}:{:.4f}'.format(itemn, mean[itemn]) \ |
| for itemn in sorted(mean.keys()) \ |
| if itemn.find('loss') == 0 |
| ] |
| console_info += mean_info |
| console_info.append('Time:{:.2f}s'.format( |
| timeit.default_timer() - self.time_check)) |
| return ' , '.join(console_info) |
|
|
| def clear(self): |
| self.sum = {} |
| self.cnt = {} |
| self.time_check = timeit.default_timer() |
|
|
| def tensorboard_close(self): |
| if self.tb is not None: |
| self.tb.close() |
|
|
| |
|
|
| def torch_to_numpy(*argv): |
| if len(argv) > 1: |
| data = list(argv) |
| else: |
| data = argv[0] |
|
|
| if isinstance(data, torch.Tensor): |
| return data.to('cpu').detach().numpy() |
|
|
| elif isinstance(data, (list, tuple)): |
| out = [] |
| for di in data: |
| out.append(torch_to_numpy(di)) |
| return out |
|
|
| elif isinstance(data, dict): |
| out = {} |
| for ni, di in data.items(): |
| out[ni] = torch_to_numpy(di) |
| return out |
| |
| else: |
| return data |
|
|
