#!/usr/bin/env python3 # -*- coding:utf-8 -*- # Copyright (c) Megvii, Inc. and its affiliates. from loguru import logger import torch from torch.nn.parallel import DistributedDataParallel as DDP from torch.utils.tensorboard import SummaryWriter from yolox.data import DataPrefetcher from yolox.utils import ( MeterBuffer, ModelEMA, all_reduce_norm, get_model_info, get_rank, get_world_size, gpu_mem_usage, load_ckpt, occupy_mem, save_checkpoint, setup_logger, synchronize ) import datetime import os import time class Trainer: def __init__(self, exp, args): # init function only defines some basic attr, other attrs like model, optimizer are built in # before_train methods. self.exp = exp self.args = args # training related attr self.max_epoch = exp.max_epoch self.amp_training = args.fp16 self.scaler = torch.cuda.amp.GradScaler(enabled=args.fp16) self.is_distributed = get_world_size() > 1 self.rank = get_rank() self.local_rank = args.local_rank self.device = "cuda:{}".format(self.local_rank) self.use_model_ema = exp.ema # data/dataloader related attr self.data_type = torch.float16 if args.fp16 else torch.float32 self.input_size = exp.input_size self.best_ap = 0 # metric record self.meter = MeterBuffer(window_size=exp.print_interval) self.file_name = os.path.join(exp.output_dir, args.experiment_name) if self.rank == 0: os.makedirs(self.file_name, exist_ok=True) setup_logger( self.file_name, distributed_rank=self.rank, filename="train_log.txt", mode="a", ) def train(self): self.before_train() try: self.train_in_epoch() except Exception: raise finally: self.after_train() def train_in_epoch(self): for self.epoch in range(self.start_epoch, self.max_epoch): self.before_epoch() self.train_in_iter() self.after_epoch() def train_in_iter(self): for self.iter in range(self.max_iter): self.before_iter() self.train_one_iter() self.after_iter() def train_one_iter(self): iter_start_time = time.time() inps, targets = self.prefetcher.next() inps = inps.to(self.data_type) targets = targets.to(self.data_type) targets.requires_grad = False data_end_time = time.time() with torch.cuda.amp.autocast(enabled=self.amp_training): outputs = self.model(inps, targets) loss = outputs["total_loss"] self.optimizer.zero_grad() self.scaler.scale(loss).backward() self.scaler.step(self.optimizer) self.scaler.update() if self.use_model_ema: self.ema_model.update(self.model) lr = self.lr_scheduler.update_lr(self.progress_in_iter + 1) for param_group in self.optimizer.param_groups: param_group["lr"] = lr iter_end_time = time.time() self.meter.update( iter_time=iter_end_time - iter_start_time, data_time=data_end_time - iter_start_time, lr=lr, **outputs, ) def before_train(self): logger.info("args: {}".format(self.args)) logger.info("exp value:\n{}".format(self.exp)) # model related init torch.cuda.set_device(self.local_rank) model = self.exp.get_model() logger.info( "Model Summary: {}".format(get_model_info(model, self.exp.test_size)) ) model.to(self.device) # solver related init self.optimizer = self.exp.get_optimizer(self.args.batch_size) # value of epoch will be set in `resume_train` model = self.resume_train(model) # data related init self.no_aug = self.start_epoch >= self.max_epoch - self.exp.no_aug_epochs self.train_loader = self.exp.get_data_loader( batch_size=self.args.batch_size, is_distributed=self.is_distributed, no_aug=self.no_aug, ) logger.info("init prefetcher, this might take one minute or less...") self.prefetcher = DataPrefetcher(self.train_loader) # max_iter means iters per epoch self.max_iter = len(self.train_loader) self.lr_scheduler = self.exp.get_lr_scheduler( self.exp.basic_lr_per_img * self.args.batch_size, self.max_iter ) if self.args.occupy: occupy_mem(self.local_rank) if self.is_distributed: model = DDP(model, device_ids=[self.local_rank], broadcast_buffers=False) if self.use_model_ema: self.ema_model = ModelEMA(model, 0.9998) self.ema_model.updates = self.max_iter * self.start_epoch self.model = model self.model.train() self.evaluator = self.exp.get_evaluator( batch_size=self.args.batch_size, is_distributed=self.is_distributed ) # Tensorboard logger if self.rank == 0: self.tblogger = SummaryWriter(self.file_name) logger.info("Training start...") #logger.info("\n{}".format(model)) def after_train(self): logger.info( "Training of experiment is done and the best AP is {:.2f}".format( self.best_ap * 100 ) ) def before_epoch(self): logger.info("---> start train epoch{}".format(self.epoch + 1)) if self.epoch + 1 == self.max_epoch - self.exp.no_aug_epochs or self.no_aug: logger.info("--->No mosaic aug now!") self.train_loader.close_mosaic() logger.info("--->Add additional L1 loss now!") if self.is_distributed: self.model.module.head.use_l1 = True else: self.model.head.use_l1 = True self.exp.eval_interval = 1 if not self.no_aug: self.save_ckpt(ckpt_name="last_mosaic_epoch") def after_epoch(self): if self.use_model_ema: self.ema_model.update_attr(self.model) self.save_ckpt(ckpt_name="latest") if (self.epoch + 1) % self.exp.eval_interval == 0: all_reduce_norm(self.model) self.evaluate_and_save_model() def before_iter(self): pass def after_iter(self): """ `after_iter` contains two parts of logic: * log information * reset setting of resize """ # log needed information if (self.iter + 1) % self.exp.print_interval == 0: # TODO check ETA logic left_iters = self.max_iter * self.max_epoch - (self.progress_in_iter + 1) eta_seconds = self.meter["iter_time"].global_avg * left_iters eta_str = "ETA: {}".format(datetime.timedelta(seconds=int(eta_seconds))) progress_str = "epoch: {}/{}, iter: {}/{}".format( self.epoch + 1, self.max_epoch, self.iter + 1, self.max_iter ) loss_meter = self.meter.get_filtered_meter("loss") loss_str = ", ".join( ["{}: {:.3f}".format(k, v.latest) for k, v in loss_meter.items()] ) time_meter = self.meter.get_filtered_meter("time") time_str = ", ".join( ["{}: {:.3f}s".format(k, v.avg) for k, v in time_meter.items()] ) logger.info( "{}, mem: {:.0f}Mb, {}, {}, lr: {:.3e}".format( progress_str, gpu_mem_usage(), time_str, loss_str, self.meter["lr"].latest, ) + (", size: {:d}, {}".format(self.input_size[0], eta_str)) ) self.meter.clear_meters() # random resizing if self.exp.random_size is not None and (self.progress_in_iter + 1) % 10 == 0: self.input_size = self.exp.random_resize( self.train_loader, self.epoch, self.rank, self.is_distributed ) @property def progress_in_iter(self): return self.epoch * self.max_iter + self.iter def resume_train(self, model): if self.args.resume: logger.info("resume training") if self.args.ckpt is None: ckpt_file = os.path.join(self.file_name, "latest" + "_ckpt.pth.tar") else: ckpt_file = self.args.ckpt ckpt = torch.load(ckpt_file, map_location=self.device) # resume the model/optimizer state dict model.load_state_dict(ckpt["model"]) self.optimizer.load_state_dict(ckpt["optimizer"]) start_epoch = ( self.args.start_epoch - 1 if self.args.start_epoch is not None else ckpt["start_epoch"] ) self.start_epoch = start_epoch logger.info( "loaded checkpoint '{}' (epoch {})".format( self.args.resume, self.start_epoch ) ) # noqa else: if self.args.ckpt is not None: logger.info("loading checkpoint for fine tuning") ckpt_file = self.args.ckpt ckpt = torch.load(ckpt_file, map_location=self.device)["model"] model = load_ckpt(model, ckpt) self.start_epoch = 0 return model def evaluate_and_save_model(self): evalmodel = self.ema_model.ema if self.use_model_ema else self.model ap50_95, ap50, summary = self.exp.eval( evalmodel, self.evaluator, self.is_distributed ) self.model.train() if self.rank == 0: self.tblogger.add_scalar("val/COCOAP50", ap50, self.epoch + 1) self.tblogger.add_scalar("val/COCOAP50_95", ap50_95, self.epoch + 1) logger.info("\n" + summary) synchronize() #self.best_ap = max(self.best_ap, ap50_95) self.save_ckpt("last_epoch", ap50 > self.best_ap) self.best_ap = max(self.best_ap, ap50) def save_ckpt(self, ckpt_name, update_best_ckpt=False): if self.rank == 0: save_model = self.ema_model.ema if self.use_model_ema else self.model logger.info("Save weights to {}".format(self.file_name)) ckpt_state = { "start_epoch": self.epoch + 1, "model": save_model.state_dict(), "optimizer": self.optimizer.state_dict(), } save_checkpoint( ckpt_state, update_best_ckpt, self.file_name, ckpt_name, )