Method __init__

trainer/base_trainer.py:30–38

(self, save_dir, log_level: str, rank: int = 0, mode='train')

    """base trainer with data parallel and tensor parallel functionalities."""

    def __init__(self, save_dir, log_level: str, rank: int = 0, mode='train') -> None:
        self.setup_save(save_dir, log_level, rank, mode=mode)
        self.rank = rank
        timer_names = [
            'data', 'vae', 'text_encoder', 'forward', 'backward', 'log',
            'ckpt', 'ema', 'image_encoder', 'controlnet'
        ]
        self.timer = EasyDict({k: Timer(k, self.logger) for k in timer_names})
        pass

    def save_config(self, config: DictConfig):
        out_path = os.path.join(self.save_dir, 'config.yaml')