base trainer with data parallel and tensor parallel functionalities.
| 25 | |
| 26 | |
| 27 | class TrainerBase: |
| 28 | """base trainer with data parallel and tensor parallel functionalities.""" |
| 29 | |
| 30 | def __init__(self, save_dir, log_level: str, rank: int = 0, mode='train') -> None: |
| 31 | self.setup_save(save_dir, log_level, rank, mode=mode) |
| 32 | self.rank = rank |
| 33 | timer_names = [ |
| 34 | 'data', 'vae', 'text_encoder', 'forward', 'backward', 'log', |
| 35 | 'ckpt', 'ema', 'image_encoder', 'controlnet' |
| 36 | ] |
| 37 | self.timer = EasyDict({k: Timer(k, self.logger) for k in timer_names}) |
| 38 | pass |
| 39 | |
| 40 | def save_config(self, config: DictConfig): |
| 41 | out_path = os.path.join(self.save_dir, 'config.yaml') |
| 42 | if self.rank == 0: |
| 43 | OmegaConf.save(config, out_path) |
| 44 | |
| 45 | def setup_save(self, save_dir: str, log_level: str, rank: int = 0, mode='train'): |
| 46 | log_level_str = log_level.upper() |
| 47 | self.save_dir = save_dir |
| 48 | self.ckpt_dir = osp.join(save_dir, 'checkpointing') |
| 49 | os.makedirs(self.ckpt_dir, exist_ok=True) |
| 50 | # self.model_weight_dir = osp.join(save_dir, 'weights') |
| 51 | # os.makedirs(self.model_weight_dir, exist_ok=True) |
| 52 | if mode == 'train': |
| 53 | self.vis_dir = osp.join(save_dir, 'visualization') |
| 54 | else: |
| 55 | self.vis_dir = osp.join(save_dir, 'test_visualization') |
| 56 | os.makedirs(self.vis_dir, exist_ok=True) |
| 57 | self.logfile = osp.join(save_dir, 'train.log') |
| 58 | if rank == 0: |
| 59 | handler = logging.FileHandler(self.logfile, mode='a') |
| 60 | handler.setLevel(log_level_str) |
| 61 | if dist.is_initialized(): |
| 62 | dist.barrier() |
| 63 | if rank != 0: |
| 64 | handler = logging.FileHandler(self.logfile, mode='a') |
| 65 | handler.setLevel(log_level_str) |
| 66 | formatter = logging.Formatter('%(asctime)s %(levelname)s %(message)s') |
| 67 | handler.setFormatter(formatter) |
| 68 | console_handler = logging.StreamHandler() |
| 69 | console_handler.setLevel(log_level_str) |
| 70 | console_handler.setFormatter(formatter) |
| 71 | logger = get_logger(__name__, log_level=log_level_str) |
| 72 | logger.logger.addHandler(handler) |
| 73 | logger.logger.addHandler(console_handler) |
| 74 | self.logger = logger |
| 75 | try: |
| 76 | self.tb_tracker = TbTracker(save_dir) |
| 77 | except ValueError: |
| 78 | pass |
| 79 | |
| 80 | # @staticmethod |
| 81 | # def load_ckpt(model, optimizer=None, ckpt_dir:str=None): |
| 82 | # state_dict = AppState(model, optimizer).state_dict() |
| 83 | # DCP.load(state_dict, checkpoint_id=ckpt_dir) |
| 84 |