import argparse import torch import torch.backends.cudnn as cudnn from config import cfg import os.path as osp # ddp import torch.distributed as dist from common.utils.distribute_utils import ( init_distributed_mode, is_main_process, set_seed ) import torch.distributed as dist from mmcv.runner import get_dist_info import math def parse_args(): parser = argparse.ArgumentParser() # parser.add_argument('--gpu', type=str, dest='gpu_ids') parser.add_argument('--num_gpus', type=int, dest='num_gpus') parser.add_argument('--master_port', type=int, dest='master_port') parser.add_argument('--exp_name', type=str, default='output/test') parser.add_argument('--config', type=str, default='./config/config_base.py') args = parser.parse_args() return args # def main(): # args = parse_args() # config_path = osp.join('./config', args.config) # cfg.get_config_fromfile(config_path) # cfg.update_config(args.num_gpus, args.exp_name) # cudnn.benchmark = True # set_seed(2023) # # ddp by default in this branch # distributed, gpu_idx = init_distributed_mode(args.master_port) # from base import Trainer # trainer = Trainer(distributed, gpu_idx) # # ddp # if distributed: # trainer.logger_info('### Set DDP ###') # trainer.logger.info(f'Distributed: {distributed}, init done {gpu_idx}') # else: # raise Exception("DDP not setup properly") # trainer.logger_info(f"Using {cfg.num_gpus} GPUs, batch size {cfg.train_batch_size} per GPU.") # trainer._make_batch_generator() # trainer._make_model() # trainer.logger_info('### Set some hyper parameters ###') # for k in cfg.__dict__: # trainer.logger_info(f'set {k} to {cfg.__dict__[k]}') # trainer.logger_info(f'train with train_3d={cfg.trainset_3d}') # trainer.logger_info(f'train with train_2d={cfg.trainset_2d}') # trainer.logger_info(f'train with trainset_humandata={cfg.trainset_humandata}') # trainer.logger_info('### Start training ###') # for epoch in range(trainer.start_epoch, cfg.end_epoch): # trainer.tot_timer.tic() # trainer.read_timer.tic() # # ddp, align random seed between devices # trainer.batch_generator.sampler.set_epoch(epoch) # for itr, (inputs, targets, meta_info) in enumerate(trainer.batch_generator): # trainer.read_timer.toc() # trainer.gpu_timer.tic() # # forward # trainer.optimizer.zero_grad() # loss= trainer.model(inputs, targets, meta_info, 'train') # loss_mean = {k: loss[k].mean() for k in loss} # loss_sum = sum(loss_mean[k] for k in loss_mean) # # backward # loss_sum.backward() # trainer.optimizer.step() # trainer.scheduler.step() # trainer.gpu_timer.toc() # if (itr + 1) % cfg.print_iters == 0: # # loss of all ranks # rank, world_size = get_dist_info() # loss_print = loss_mean.copy() # for k in loss_print: # dist.all_reduce(loss_print[k]) # total_loss = 0 # for k in loss_print: # loss_print[k] = loss_print[k] / world_size # total_loss += loss_print[k] # loss_print['total'] = total_loss # screen = [ # 'Epoch %d/%d itr %d/%d:' % (epoch, cfg.end_epoch, itr, trainer.itr_per_epoch), # 'lr: %g' % (trainer.get_lr()), # 'speed: %.2f(%.2fs r%.2f)s/itr' % ( # trainer.tot_timer.average_time, trainer.gpu_timer.average_time, # trainer.read_timer.average_time), # '%.2fh/epoch' % (trainer.tot_timer.average_time / 3600. * trainer.itr_per_epoch), # ] # screen += ['%s: %.4f' % ('loss_' + k, v.detach()) for k, v in loss_print.items()] # trainer.logger_info(' '.join(screen)) # trainer.tot_timer.toc() # trainer.tot_timer.tic() # trainer.read_timer.tic() # # save model ddp, save model.module on rank 0 only # save_epoch = getattr(cfg, 'save_epoch', 10) # if is_main_process() and (epoch % save_epoch == 0 or epoch == cfg.end_epoch - 1): # trainer.save_model({ # 'epoch': epoch, # 'network': trainer.model.state_dict(), # 'optimizer': trainer.optimizer.state_dict(), # }, epoch) # dist.barrier() def main(): args = parse_args() config_path = osp.join('./config', args.config) cfg.get_config_fromfile(config_path) cfg.update_config(args.num_gpus, args.exp_name) cudnn.benchmark = True set_seed(2023) # check for DDP usage distributed = False gpu_idx = 0 use_lora = cfg.use_lora if args.num_gpus > 1: distributed, gpu_idx = init_distributed_mode(args.master_port) else: print("Running in single GPU mode.") from base import Trainer trainer = Trainer(distributed, gpu_idx, use_lora) trainer.logger_info(f"Using {cfg.num_gpus} GPU(s), batch size {cfg.train_batch_size} per GPU.") trainer._make_batch_generator() trainer._make_model() trainer.logger_info('### Set some hyper parameters ###') for k in cfg.__dict__: trainer.logger_info(f'set {k} to {cfg.__dict__[k]}') trainer.logger_info(f'train with train_3d={cfg.trainset_3d}') trainer.logger_info(f'train with train_2d={cfg.trainset_2d}') trainer.logger_info(f'train with trainset_humandata={cfg.trainset_humandata}') trainer.logger_info('### Start training ###') for epoch in range(trainer.start_epoch, cfg.end_epoch): trainer.tot_timer.tic() trainer.read_timer.tic() # set_epoch only when using distributed sampler if distributed: trainer.batch_generator.sampler.set_epoch(epoch) for itr, (inputs, targets, meta_info) in enumerate(trainer.batch_generator): trainer.read_timer.toc() trainer.gpu_timer.tic() trainer.optimizer.zero_grad() # loss = trainer.model(inputs, targets, meta_info, 'train') # loss_mean = {k: loss[k].mean() for k in loss} # Dposerx cosine weight loss = trainer.model(inputs, targets, meta_info, 'train') loss_mean = {k: loss[k].mean() for k in loss} w0 = getattr(cfg, 'dposer_x_weight', 1.0) w_min = getattr(cfg, 'dposer_x_weight_min', 0.001) warmup = getattr(cfg, 'dposer_x_weight_warmup', 2) if epoch < warmup: w_dposerx = w0 else: frac = (itr + 1) / max(1, trainer.itr_per_epoch) t = (epoch + frac - warmup) / max(1, cfg.end_epoch - warmup) t = min(max(t, 0.0), 1.0) w_dposerx = w_min + 0.5 * (w0 - w_min) * (1.0 + math.cos(math.pi * t)) if 'dposerx' in loss_mean: loss_mean['dposerx'] = loss_mean['dposerx'] * w_dposerx loss_sum = sum(loss_mean[k] for k in loss_mean) loss_sum.backward() trainer.optimizer.step() trainer.scheduler.step() trainer.gpu_timer.toc() if (itr + 1) % cfg.print_iters == 0: if distributed: rank, world_size = get_dist_info() loss_print = loss_mean.copy() for k in loss_print: dist.all_reduce(loss_print[k]) total_loss = 0 for k in loss_print: loss_print[k] = loss_print[k] / world_size total_loss += loss_print[k] loss_print['total'] = total_loss else: loss_print = loss_mean loss_print['total'] = loss_sum screen = [ 'Epoch %d/%d itr %d/%d:' % (epoch, cfg.end_epoch, itr, trainer.itr_per_epoch), 'lr: %g' % (trainer.get_lr()), 'speed: %.2f(%.2fs r%.2f)s/itr' % ( trainer.tot_timer.average_time, trainer.gpu_timer.average_time, trainer.read_timer.average_time), '%.2fh/epoch' % (trainer.tot_timer.average_time / 3600. * trainer.itr_per_epoch), ] screen += ['%s: %.4f' % ('loss_' + k, v.detach()) for k, v in loss_print.items()] screen += [f'dposerx_w: {float(w_dposerx):.4f}'] trainer.logger_info(' '.join(screen)) trainer.tot_timer.toc() trainer.tot_timer.tic() trainer.read_timer.tic() # save model if not distributed or is_main_process(): save_epoch = getattr(cfg, 'save_epoch', 10) if epoch % save_epoch == 0 or epoch == cfg.end_epoch - 1: trainer.save_model({ 'epoch': epoch, 'network': trainer.model.state_dict(), 'optimizer': trainer.optimizer.state_dict(), }, epoch) if distributed: dist.barrier() if __name__ == "__main__": main()