def setup_after_launch(
cfg: CfgNode,
output_dir: str,
runner: Optional[BaseRunner] = None,
_scale_world_size:
bool = True, # HACK: temporarily allow lightning_train_net to by pass this.
):
"""
Binary-level setup after entering DDP, including
- creating working directory
- setting up logger
- logging environment
- printing and dumping config
- (optional) initializing runner
"""
create_dir_on_global_main_process(output_dir)
setup_loggers(output_dir)
log_system_info()
cfg.freeze()
maybe_override_output_dir(cfg, output_dir)
logger.info("Running with full config:\n{}".format(cfg))
dump_cfg(cfg, os.path.join(output_dir, "config.yaml"))
if runner:
logger.info("Initializing runner ...")
runner = initialize_runner(runner, cfg)
logger.info("Running with runner: {}".format(runner))
# save the diff config
if runner:
default_cfg = runner.get_default_cfg()
dump_cfg(
get_diff_cfg(default_cfg, cfg),
os.path.join(output_dir, "diff_config.yaml"),
)
else:
# TODO: support getting default_cfg without runner.
pass
# scale the config after dumping so that dumped config files keep original world size
if _scale_world_size:
auto_scale_world_size(cfg, new_world_size=comm.get_world_size())
def _setup_after_launch(cfg: CN, output_dir: str, runner):
"""
Set things up after entering DDP, including
- creating working directory
- setting up logger
- logging environment
- initializing runner
"""
create_dir_on_global_main_process(output_dir)
comm.synchronize()
setup_loggers(output_dir)
cfg.freeze()
if cfg.OUTPUT_DIR != output_dir:
with temp_defrost(cfg):
logger.warning(
"Override cfg.OUTPUT_DIR ({}) to be the same as output_dir {}".format(
cfg.OUTPUT_DIR, output_dir
)
)
cfg.OUTPUT_DIR = output_dir
dump_cfg(cfg, os.path.join(output_dir, "config.yaml"))