Spaces:
Runtime error
Runtime error
File size: 3,107 Bytes
0328207 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 | import argparse
import os
import time
import hydra
from hydra.core.hydra_config import HydraConfig
from omegaconf import DictConfig, OmegaConf
from trainer import (
DiffusionTrainer,
GANTrainer,
ODETrainer,
ScoreDistillationTrainer,
)
import torch.distributed as dist
import logging
from pathlib import Path
from utils.misc import format_dict
def _current_node_rank():
if dist.is_available() and dist.is_initialized():
return dist.get_rank() // dist.get_world_size()
# fallback to env vars set by torchrun/launch
return int(os.environ.get("NODE_RANK", 0))
def _current_process_rank():
if dist.is_available() and dist.is_initialized():
return dist.get_rank()
# fallback to env vars set by torchrun/launch
return int(os.environ.get("RANK", os.environ.get("LOCAL_RANK", 0)))
def _add_rank_to_record():
factory = logging.getLogRecordFactory()
def record_factory(*args, **kwargs):
record = factory(*args, **kwargs)
record.rank = _current_process_rank()
return record
logging.setLogRecordFactory(record_factory)
def _configure_logging(
logdir: Path,
):
"""
Hydra installs its own logging handlers before main() runs.
This function forcefully replaces them so our file + format take effect.
"""
_add_rank_to_record()
ts = time.strftime("%Y-%m-%d_%H-%M-%S")
logfile = logdir / f"train-node{_current_node_rank()}-rank{_current_process_rank()}-{ts}.log"
fmt = logging.Formatter(
"[rank:{rank}] [{levelname}] [{asctime}] : {message}",
style="{",
)
file_handler = logging.FileHandler(logfile, mode="a")
file_handler.setLevel(logging.INFO)
file_handler.setFormatter(fmt)
stream_handler = logging.StreamHandler()
stream_handler.setLevel(logging.INFO)
stream_handler.setFormatter(fmt)
# Replace existing handlers (Hydra already configured them)
root = logging.getLogger()
root.setLevel(logging.INFO)
root.handlers[:] = [file_handler, stream_handler]
@hydra.main(
version_base=None,
config_path="configs",
config_name="rolling_forcing_dmd",
)
def main(
config: DictConfig,
):
# Strict mode: CLI override keys must already exist in cfg, otherwise error.
OmegaConf.set_struct(config, True)
config_name = HydraConfig.get().job.config_name
# Keep paths stable even if Hydra changes working dir
orig_cwd = hydra.utils.get_original_cwd()
logdir = Path(config.logdir)
if not logdir.is_absolute():
logdir = Path(orig_cwd) / logdir
logdir.mkdir(parents=True, exist_ok=True)
_configure_logging(logdir)
logging.info(f"""
{config_name = }
config = {format_dict(config)}
""")
if config.trainer == "diffusion":
trainer = DiffusionTrainer(config)
elif config.trainer == "gan":
trainer = GANTrainer(config)
elif config.trainer == "ode":
trainer = ODETrainer(config)
elif config.trainer == "score_distillation":
trainer = ScoreDistillationTrainer(config)
trainer.train()
if __name__ == "__main__":
main()
|