Spaces:
Running
Running
| import copy | |
| import datetime | |
| import os | |
| import random | |
| import time | |
| import numpy as np | |
| import torch | |
| from tqdm import tqdm | |
| from openrec.losses import build_loss | |
| from openrec.metrics import build_metric | |
| from openrec.modeling import build_model | |
| from openrec.optimizer import build_optimizer | |
| from openrec.postprocess import build_post_process | |
| from tools.data import build_dataloader | |
| from tools.utils.ckpt import load_ckpt, save_ckpt | |
| from tools.utils.logging import get_logger | |
| from tools.utils.stats import TrainingStats | |
| from tools.utils.utility import AverageMeter | |
| __all__ = ['Trainer'] | |
| def get_parameter_number(model): | |
| total_num = sum(p.numel() for p in model.parameters()) | |
| trainable_num = sum(p.numel() for p in model.parameters() | |
| if p.requires_grad) | |
| return {'Total': total_num, 'Trainable': trainable_num} | |
| class Trainer(object): | |
| def __init__(self, cfg, mode='train'): | |
| self.cfg = cfg.cfg | |
| self.local_rank = (int(os.environ['LOCAL_RANK']) | |
| if 'LOCAL_RANK' in os.environ else 0) | |
| self.set_device(self.cfg['Global']['device']) | |
| mode = mode.lower() | |
| assert mode in [ | |
| 'train_eval', | |
| 'train', | |
| 'eval', | |
| 'test', | |
| ], 'mode should be train, eval and test' | |
| if torch.cuda.device_count() > 1 and 'train' in mode: | |
| torch.distributed.init_process_group(backend='nccl') | |
| torch.cuda.set_device(self.device) | |
| self.cfg['Global']['distributed'] = True | |
| else: | |
| self.cfg['Global']['distributed'] = False | |
| self.local_rank = 0 | |
| self.cfg['Global']['output_dir'] = self.cfg['Global'].get( | |
| 'output_dir', 'output') | |
| os.makedirs(self.cfg['Global']['output_dir'], exist_ok=True) | |
| self.writer = None | |
| if self.local_rank == 0 and self.cfg['Global'][ | |
| 'use_tensorboard'] and 'train' in mode: | |
| from torch.utils.tensorboard import SummaryWriter | |
| self.writer = SummaryWriter(self.cfg['Global']['output_dir']) | |
| self.logger = get_logger( | |
| 'openrec', | |
| os.path.join(self.cfg['Global']['output_dir'], 'train.log') | |
| if 'train' in mode else None, | |
| ) | |
| cfg.print_cfg(self.logger.info) | |
| if self.cfg['Global']['device'] == 'gpu' and self.device.type == 'cpu': | |
| self.logger.info('cuda is not available, auto switch to cpu') | |
| self.grad_clip_val = self.cfg['Global'].get('grad_clip_val', 0) | |
| self.all_ema = self.cfg['Global'].get('all_ema', True) | |
| self.use_ema = self.cfg['Global'].get('use_ema', True) | |
| self.set_random_seed(self.cfg['Global'].get('seed', 48)) | |
| # build data loader | |
| self.train_dataloader = None | |
| if 'train' in mode: | |
| cfg.save( | |
| os.path.join(self.cfg['Global']['output_dir'], 'config.yml'), | |
| self.cfg) | |
| self.train_dataloader = build_dataloader(self.cfg, 'Train', | |
| self.logger) | |
| self.logger.info( | |
| f'train dataloader has {len(self.train_dataloader)} iters') | |
| self.valid_dataloader = None | |
| if 'eval' in mode and self.cfg['Eval']: | |
| self.valid_dataloader = build_dataloader(self.cfg, 'Eval', | |
| self.logger) | |
| self.logger.info( | |
| f'valid dataloader has {len(self.valid_dataloader)} iters') | |
| # build post process | |
| self.post_process_class = build_post_process(self.cfg['PostProcess'], | |
| self.cfg['Global']) | |
| # build model | |
| # for rec algorithm | |
| char_num = self.post_process_class.get_character_num() | |
| self.cfg['Architecture']['Decoder']['out_channels'] = char_num | |
| self.model = build_model(self.cfg['Architecture']) | |
| self.logger.info(get_parameter_number(model=self.model)) | |
| self.model = self.model.to(self.device) | |
| if self.local_rank == 0: | |
| ema_model = build_model(self.cfg['Architecture']) | |
| self.ema_model = ema_model.to(self.device) | |
| self.ema_model.eval() | |
| use_sync_bn = self.cfg['Global'].get('use_sync_bn', False) | |
| if use_sync_bn: | |
| self.model = torch.nn.SyncBatchNorm.convert_sync_batchnorm( | |
| self.model) | |
| self.logger.info('convert_sync_batchnorm') | |
| # build loss | |
| self.loss_class = build_loss(self.cfg['Loss']) | |
| self.optimizer, self.lr_scheduler = None, None | |
| if self.train_dataloader is not None: | |
| # build optim | |
| self.optimizer, self.lr_scheduler = build_optimizer( | |
| self.cfg['Optimizer'], | |
| self.cfg['LRScheduler'], | |
| epochs=self.cfg['Global']['epoch_num'], | |
| step_each_epoch=len(self.train_dataloader), | |
| model=self.model, | |
| ) | |
| self.eval_class = build_metric(self.cfg['Metric']) | |
| self.status = load_ckpt(self.model, self.cfg, self.optimizer, | |
| self.lr_scheduler) | |
| if self.cfg['Global']['distributed']: | |
| self.model = torch.nn.parallel.DistributedDataParallel( | |
| self.model, [self.local_rank], find_unused_parameters=False) | |
| # amp | |
| self.scaler = (torch.cuda.amp.GradScaler() if self.cfg['Global'].get( | |
| 'use_amp', False) else None) | |
| self.logger.info( | |
| f'run with torch {torch.__version__} and device {self.device}') | |
| def load_params(self, params): | |
| self.model.load_state_dict(params) | |
| def set_random_seed(self, seed): | |
| torch.manual_seed(seed) # 为CPU设置随机种子 | |
| if self.device.type == 'cuda': | |
| torch.backends.cudnn.benchmark = True | |
| torch.cuda.manual_seed(seed) # 为当前GPU设置随机种子 | |
| torch.cuda.manual_seed_all(seed) # 为所有GPU设置随机种子 | |
| random.seed(seed) | |
| np.random.seed(seed) | |
| def set_device(self, device): | |
| if device == 'gpu' and torch.cuda.is_available(): | |
| device = torch.device(f'cuda:{self.local_rank}') | |
| else: | |
| device = torch.device('cpu') | |
| self.device = device | |
| def train(self): | |
| cal_metric_during_train = self.cfg['Global'].get( | |
| 'cal_metric_during_train', False) | |
| log_smooth_window = self.cfg['Global']['log_smooth_window'] | |
| epoch_num = self.cfg['Global']['epoch_num'] | |
| print_batch_step = self.cfg['Global']['print_batch_step'] | |
| eval_epoch_step = self.cfg['Global'].get('eval_epoch_step', 1) | |
| start_eval_epoch = 0 | |
| if self.valid_dataloader is not None: | |
| if type(eval_epoch_step) == list and len(eval_epoch_step) >= 2: | |
| start_eval_epoch = eval_epoch_step[0] | |
| eval_epoch_step = eval_epoch_step[1] | |
| if len(self.valid_dataloader) == 0: | |
| start_eval_epoch = 1e111 | |
| self.logger.info( | |
| 'No Images in eval dataset, evaluation during training will be disabled' | |
| ) | |
| self.logger.info( | |
| f'During the training process, after the {start_eval_epoch}th epoch, ' | |
| f'an evaluation is run every {eval_epoch_step} epoch') | |
| else: | |
| start_eval_epoch = 1e111 | |
| eval_batch_step = self.cfg['Global']['eval_batch_step'] | |
| global_step = self.status.get('global_step', 0) | |
| start_eval_step = 0 | |
| if type(eval_batch_step) == list and len(eval_batch_step) >= 2: | |
| start_eval_step = eval_batch_step[0] | |
| eval_batch_step = eval_batch_step[1] | |
| if len(self.valid_dataloader) == 0: | |
| self.logger.info( | |
| 'No Images in eval dataset, evaluation during training ' | |
| 'will be disabled') | |
| start_eval_step = 1e111 | |
| self.logger.info( | |
| 'During the training process, after the {}th iteration, ' | |
| 'an evaluation is run every {} iterations'.format( | |
| start_eval_step, eval_batch_step)) | |
| start_epoch = self.status.get('epoch', 1) | |
| best_metric = self.status.get('metrics', {}) | |
| if self.eval_class.main_indicator not in best_metric: | |
| best_metric[self.eval_class.main_indicator] = 0 | |
| ema_best_metric = self.status.get('metrics', {}) | |
| ema_best_metric[self.eval_class.main_indicator] = 0 | |
| train_stats = TrainingStats(log_smooth_window, ['lr']) | |
| self.model.train() | |
| total_samples = 0 | |
| train_reader_cost = 0.0 | |
| train_batch_cost = 0.0 | |
| best_iter = 0 | |
| ema_stpe = 1 | |
| ema_eval_iter = 0 | |
| loss_avg = 0. | |
| reader_start = time.time() | |
| eta_meter = AverageMeter() | |
| for epoch in range(start_epoch, epoch_num + 1): | |
| if self.train_dataloader.dataset.need_reset: | |
| self.train_dataloader = build_dataloader( | |
| self.cfg, | |
| 'Train', | |
| self.logger, | |
| epoch=epoch % 20 if epoch % 20 != 0 else 20, | |
| ) | |
| for idx, batch in enumerate(self.train_dataloader): | |
| batch = [t.to(self.device) for t in batch] | |
| self.optimizer.zero_grad() | |
| train_reader_cost += time.time() - reader_start | |
| # use amp | |
| if self.scaler: | |
| with torch.cuda.amp.autocast(): | |
| preds = self.model(batch[0], data=batch[1:]) | |
| loss = self.loss_class(preds, batch) | |
| self.scaler.scale(loss['loss']).backward() | |
| if self.grad_clip_val > 0: | |
| torch.nn.utils.clip_grad_norm_( | |
| self.model.parameters(), | |
| max_norm=self.grad_clip_val) | |
| self.scaler.step(self.optimizer) | |
| self.scaler.update() | |
| else: | |
| preds = self.model(batch[0], data=batch[1:]) | |
| loss = self.loss_class(preds, batch) | |
| avg_loss = loss['loss'] | |
| avg_loss.backward() | |
| if self.grad_clip_val > 0: | |
| torch.nn.utils.clip_grad_norm_( | |
| self.model.parameters(), | |
| max_norm=self.grad_clip_val) | |
| self.optimizer.step() | |
| if cal_metric_during_train: # only rec and cls need | |
| post_result = self.post_process_class(preds, | |
| batch, | |
| training=True) | |
| self.eval_class(post_result, batch, training=True) | |
| metric = self.eval_class.get_metric() | |
| train_stats.update(metric) | |
| train_batch_time = time.time() - reader_start | |
| train_batch_cost += train_batch_time | |
| eta_meter.update(train_batch_time) | |
| global_step += 1 | |
| total_samples += len(batch[0]) | |
| self.lr_scheduler.step() | |
| if self.local_rank == 0 and self.use_ema and epoch > ( | |
| epoch_num - epoch_num // 10): | |
| with torch.no_grad(): | |
| loss_currn = loss['loss'].detach().cpu().numpy().mean() | |
| loss_avg = ((loss_avg * | |
| (ema_stpe - 1)) + loss_currn) / (ema_stpe) | |
| if ema_stpe == 1: | |
| # current_weight = copy.deepcopy(self.model.module.state_dict()) | |
| ema_state_dict = copy.deepcopy( | |
| self.model.module.state_dict() if self. | |
| cfg['Global']['distributed'] else self.model. | |
| state_dict()) | |
| self.ema_model.load_state_dict(ema_state_dict) | |
| # if global_step > (epoch_num - epoch_num//10)*max_iter: | |
| elif loss_currn <= loss_avg or self.all_ema: | |
| # eval_batch_step = 500 | |
| current_weight = copy.deepcopy( | |
| self.model.module.state_dict() if self. | |
| cfg['Global']['distributed'] else self.model. | |
| state_dict()) | |
| k1 = 1 / (ema_stpe + 1) | |
| k2 = 1 - k1 | |
| for k, v in ema_state_dict.items(): | |
| # v = (v * (ema_stpe - 1) + current_weight[k])/ema_stpe | |
| v = v * k2 + current_weight[k] * k1 | |
| # v.req = True | |
| ema_state_dict[k] = v | |
| # ema_stpe += 1 | |
| self.ema_model.load_state_dict(ema_state_dict) | |
| ema_stpe += 1 | |
| if global_step > start_eval_step and ( | |
| global_step - | |
| start_eval_step) % eval_batch_step == 0: | |
| ema_cur_metric = self.eval_ema() | |
| ema_cur_metric_str = f"cur ema metric, {', '.join(['{}: {}'.format(k, v) for k, v in ema_cur_metric.items()])}" | |
| self.logger.info(ema_cur_metric_str) | |
| state = { | |
| 'epoch': epoch, | |
| 'global_step': global_step, | |
| 'state_dict': self.ema_model.state_dict(), | |
| 'optimizer': None, | |
| 'scheduler': None, | |
| 'config': self.cfg, | |
| 'metrics': ema_cur_metric, | |
| } | |
| save_path = os.path.join( | |
| self.cfg['Global']['output_dir'], | |
| 'ema_' + str(ema_eval_iter) + '.pth') | |
| torch.save(state, save_path) | |
| self.logger.info(f'save ema ckpt to {save_path}') | |
| ema_eval_iter += 1 | |
| if ema_cur_metric[self.eval_class. | |
| main_indicator] >= ema_best_metric[ | |
| self.eval_class.main_indicator]: | |
| ema_best_metric.update(ema_cur_metric) | |
| ema_best_metric['best_epoch'] = epoch | |
| best_ema_str = f"best metric, {', '.join(['{}: {}'.format(k, v) for k, v in ema_best_metric.items()])}" | |
| self.logger.info(best_ema_str) | |
| # logger | |
| stats = { | |
| k: float(v) | |
| if v.shape == [] else v.detach().cpu().numpy().mean() | |
| for k, v in loss.items() | |
| } | |
| stats['lr'] = self.lr_scheduler.get_last_lr()[0] | |
| train_stats.update(stats) | |
| if self.writer is not None: | |
| for k, v in train_stats.get().items(): | |
| self.writer.add_scalar(f'TRAIN/{k}', v, global_step) | |
| if self.local_rank == 0 and ( | |
| (global_step > 0 and global_step % print_batch_step == 0) | |
| or (idx >= len(self.train_dataloader) - 1)): | |
| logs = train_stats.log() | |
| eta_sec = ( | |
| (epoch_num + 1 - epoch) * len(self.train_dataloader) - | |
| idx - 1) * eta_meter.avg | |
| eta_sec_format = str( | |
| datetime.timedelta(seconds=int(eta_sec))) | |
| strs = ( | |
| f'epoch: [{epoch}/{epoch_num}], global_step: {global_step}, {logs}, ' | |
| f'avg_reader_cost: {train_reader_cost / print_batch_step:.5f} s, ' | |
| f'avg_batch_cost: {train_batch_cost / print_batch_step:.5f} s, ' | |
| f'avg_samples: {total_samples / print_batch_step}, ' | |
| f'ips: {total_samples / train_batch_cost:.5f} samples/s, ' | |
| f'eta: {eta_sec_format}') | |
| self.logger.info(strs) | |
| total_samples = 0 | |
| train_reader_cost = 0.0 | |
| train_batch_cost = 0.0 | |
| reader_start = time.time() | |
| # eval | |
| if (global_step > start_eval_step and | |
| (global_step - start_eval_step) % eval_batch_step | |
| == 0) and self.local_rank == 0: | |
| cur_metric = self.eval() | |
| cur_metric_str = f"cur metric, {', '.join(['{}: {}'.format(k, v) for k, v in cur_metric.items()])}" | |
| self.logger.info(cur_metric_str) | |
| # logger metric | |
| if self.writer is not None: | |
| for k, v in cur_metric.items(): | |
| if isinstance(v, (float, int)): | |
| self.writer.add_scalar(f'EVAL/{k}', | |
| cur_metric[k], | |
| global_step) | |
| if (cur_metric[self.eval_class.main_indicator] >= | |
| best_metric[self.eval_class.main_indicator]): | |
| best_metric.update(cur_metric) | |
| best_metric['best_epoch'] = epoch | |
| if self.writer is not None: | |
| self.writer.add_scalar( | |
| f'EVAL/best_{self.eval_class.main_indicator}', | |
| best_metric[self.eval_class.main_indicator], | |
| global_step, | |
| ) | |
| if epoch > (epoch_num - epoch_num // 10 - 2): | |
| save_ckpt(self.model, | |
| self.cfg, | |
| self.optimizer, | |
| self.lr_scheduler, | |
| epoch, | |
| global_step, | |
| best_metric, | |
| is_best=True, | |
| prefix='best_' + str(best_iter)) | |
| best_iter += 1 | |
| # else: | |
| save_ckpt(self.model, | |
| self.cfg, | |
| self.optimizer, | |
| self.lr_scheduler, | |
| epoch, | |
| global_step, | |
| best_metric, | |
| is_best=True, | |
| prefix=None) | |
| best_str = f"best metric, {', '.join(['{}: {}'.format(k, v) for k, v in best_metric.items()])}" | |
| self.logger.info(best_str) | |
| if self.local_rank == 0 and epoch > start_eval_epoch and ( | |
| epoch - start_eval_epoch) % eval_epoch_step == 0: | |
| cur_metric = self.eval() | |
| cur_metric_str = f"cur metric, {', '.join(['{}: {}'.format(k, v) for k, v in cur_metric.items()])}" | |
| self.logger.info(cur_metric_str) | |
| # logger metric | |
| if self.writer is not None: | |
| for k, v in cur_metric.items(): | |
| if isinstance(v, (float, int)): | |
| self.writer.add_scalar(f'EVAL/{k}', cur_metric[k], | |
| global_step) | |
| if (cur_metric[self.eval_class.main_indicator] >= | |
| best_metric[self.eval_class.main_indicator]): | |
| best_metric.update(cur_metric) | |
| best_metric['best_epoch'] = epoch | |
| if self.writer is not None: | |
| self.writer.add_scalar( | |
| f'EVAL/best_{self.eval_class.main_indicator}', | |
| best_metric[self.eval_class.main_indicator], | |
| global_step, | |
| ) | |
| if epoch > (epoch_num - epoch_num // 10 - 2): | |
| save_ckpt(self.model, | |
| self.cfg, | |
| self.optimizer, | |
| self.lr_scheduler, | |
| epoch, | |
| global_step, | |
| best_metric, | |
| is_best=True, | |
| prefix='best_' + str(best_iter)) | |
| best_iter += 1 | |
| # else: | |
| save_ckpt(self.model, | |
| self.cfg, | |
| self.optimizer, | |
| self.lr_scheduler, | |
| epoch, | |
| global_step, | |
| best_metric, | |
| is_best=True, | |
| prefix=None) | |
| best_str = f"best metric, {', '.join(['{}: {}'.format(k, v) for k, v in best_metric.items()])}" | |
| self.logger.info(best_str) | |
| if self.local_rank == 0: | |
| save_ckpt(self.model, | |
| self.cfg, | |
| self.optimizer, | |
| self.lr_scheduler, | |
| epoch, | |
| global_step, | |
| best_metric, | |
| is_best=False, | |
| prefix=None) | |
| if epoch > (epoch_num - epoch_num // 10 - 2): | |
| save_ckpt(self.model, | |
| self.cfg, | |
| self.optimizer, | |
| self.lr_scheduler, | |
| epoch, | |
| global_step, | |
| best_metric, | |
| is_best=False, | |
| prefix='epoch_' + str(epoch)) | |
| if self.use_ema and epoch > (epoch_num - epoch_num // 10): | |
| # if global_step > start_eval_step and (global_step - start_eval_step) % eval_batch_step == 0: | |
| ema_cur_metric = self.eval_ema() | |
| ema_cur_metric_str = f"cur ema metric, {', '.join(['{}: {}'.format(k, v) for k, v in ema_cur_metric.items()])}" | |
| self.logger.info(ema_cur_metric_str) | |
| state = { | |
| 'epoch': epoch, | |
| 'global_step': global_step, | |
| 'state_dict': self.ema_model.state_dict(), | |
| 'optimizer': None, | |
| 'scheduler': None, | |
| 'config': self.cfg, | |
| 'metrics': ema_cur_metric, | |
| } | |
| save_path = os.path.join( | |
| self.cfg['Global']['output_dir'], | |
| 'ema_' + str(ema_eval_iter) + '.pth') | |
| torch.save(state, save_path) | |
| self.logger.info(f'save ema ckpt to {save_path}') | |
| ema_eval_iter += 1 | |
| if (ema_cur_metric[self.eval_class.main_indicator] >= | |
| ema_best_metric[self.eval_class.main_indicator]): | |
| ema_best_metric.update(ema_cur_metric) | |
| ema_best_metric['best_epoch'] = epoch | |
| # ema_cur_metric_str = f"best ema metric, {', '.join(['{}: {}'.format(k, v) for k, v in ema_best_metric.items()])}" | |
| best_ema_str = f"best metric, {', '.join(['{}: {}'.format(k, v) for k, v in ema_best_metric.items()])}" | |
| self.logger.info(best_ema_str) | |
| best_str = f"best metric, {', '.join(['{}: {}'.format(k, v) for k, v in best_metric.items()])}" | |
| self.logger.info(best_str) | |
| if self.writer is not None: | |
| self.writer.close() | |
| if torch.cuda.device_count() > 1: | |
| torch.distributed.destroy_process_group() | |
| def eval(self): | |
| self.model.eval() | |
| with torch.no_grad(): | |
| total_frame = 0.0 | |
| total_time = 0.0 | |
| pbar = tqdm( | |
| total=len(self.valid_dataloader), | |
| desc='eval model:', | |
| position=0, | |
| leave=True, | |
| ) | |
| sum_images = 0 | |
| for idx, batch in enumerate(self.valid_dataloader): | |
| batch = [t.to(self.device) for t in batch] | |
| start = time.time() | |
| if self.scaler: | |
| with torch.cuda.amp.autocast(): | |
| preds = self.model(batch[0], data=batch[1:]) | |
| else: | |
| preds = self.model(batch[0], data=batch[1:]) | |
| total_time += time.time() - start | |
| # Obtain usable results from post-processing methods | |
| # Evaluate the results of the current batch | |
| post_result = self.post_process_class(preds, batch) | |
| self.eval_class(post_result, batch) | |
| pbar.update(1) | |
| total_frame += len(batch[0]) | |
| sum_images += 1 | |
| # Get final metric，eg. acc or hmean | |
| metric = self.eval_class.get_metric() | |
| pbar.close() | |
| self.model.train() | |
| metric['fps'] = total_frame / total_time | |
| return metric | |
| def eval_ema(self): | |
| # self.model.eval() | |
| with torch.no_grad(): | |
| total_frame = 0.0 | |
| total_time = 0.0 | |
| pbar = tqdm( | |
| total=len(self.valid_dataloader), | |
| desc='eval ema_model:', | |
| position=0, | |
| leave=True, | |
| ) | |
| sum_images = 0 | |
| for idx, batch in enumerate(self.valid_dataloader): | |
| batch = [t.to(self.device) for t in batch] | |
| start = time.time() | |
| if self.scaler: | |
| with torch.cuda.amp.autocast(): | |
| preds = self.ema_model(batch[0], data=batch[1:]) | |
| else: | |
| preds = self.ema_model(batch[0], data=batch[1:]) | |
| total_time += time.time() - start | |
| # Obtain usable results from post-processing methods | |
| # Evaluate the results of the current batch | |
| post_result = self.post_process_class(preds, batch) | |
| self.eval_class(post_result, batch) | |
| pbar.update(1) | |
| total_frame += len(batch[0]) | |
| sum_images += 1 | |
| # Get final metric，eg. acc or hmean | |
| metric = self.eval_class.get_metric() | |
| pbar.close() | |
| # self.model.train() | |
| metric['fps'] = total_frame / total_time | |
| return metric | |
| def test_dataloader(self): | |
| starttime = time.time() | |
| count = 0 | |
| try: | |
| for data in self.train_dataloader: | |
| count += 1 | |
| if count % 1 == 0: | |
| batch_time = time.time() - starttime | |
| starttime = time.time() | |
| self.logger.info( | |
| f'reader: {count}, {data[0].shape}, {batch_time}') | |
| except: | |
| import traceback | |
| self.logger.info(traceback.format_exc()) | |
| self.logger.info(f'finish reader: {count}, Success!') | |