Engine for training and testing.
| 396 | |
| 397 | |
| 398 | class Runner: |
| 399 | """Engine for training and testing.""" |
| 400 | |
| 401 | def __init__( |
| 402 | self, local_rank: int, world_rank, world_size: int, cfg: Config |
| 403 | ) -> None: |
| 404 | set_random_seed(42 + local_rank) |
| 405 | |
| 406 | self.cfg = cfg |
| 407 | self.world_rank = world_rank |
| 408 | self.local_rank = local_rank |
| 409 | self.world_size = world_size |
| 410 | self.device = f"cuda:{local_rank}" |
| 411 | |
| 412 | # Where to dump results. |
| 413 | os.makedirs(cfg.result_dir, exist_ok=True) |
| 414 | |
| 415 | # Setup output directories. |
| 416 | self.ckpt_dir = f"{cfg.result_dir}/ckpts" |
| 417 | os.makedirs(self.ckpt_dir, exist_ok=True) |
| 418 | self.stats_dir = f"{cfg.result_dir}/stats" |
| 419 | os.makedirs(self.stats_dir, exist_ok=True) |
| 420 | self.render_dir = f"{cfg.result_dir}/renders" |
| 421 | os.makedirs(self.render_dir, exist_ok=True) |
| 422 | |
| 423 | # Tensorboard |
| 424 | self.writer = SummaryWriter(log_dir=f"{cfg.result_dir}/tb") |
| 425 | |
| 426 | # Load data: Training data should contain initial points and colors. |
| 427 | self.parser = Parser( |
| 428 | data_dir=cfg.data_dir, |
| 429 | factor=cfg.data_factor, |
| 430 | normalize=cfg.normalize_world_space, |
| 431 | test_every=cfg.test_every, |
| 432 | ) |
| 433 | if cfg.scene_type == "GSC" and cfg.test_view_id is not None: # GSC mode |
| 434 | self.trainset = GSCDataset( |
| 435 | self.parser, |
| 436 | split="train", |
| 437 | patch_size=cfg.patch_size, |
| 438 | load_depths=cfg.depth_loss, |
| 439 | test_view_ids=cfg.test_view_id, |
| 440 | ) |
| 441 | self.valset = GSCDataset( |
| 442 | self.parser, |
| 443 | split="val", |
| 444 | test_view_ids=cfg.test_view_id,) |
| 445 | else: # default mode |
| 446 | self.trainset = Dataset( |
| 447 | self.parser, |
| 448 | split="train", |
| 449 | patch_size=cfg.patch_size, |
| 450 | load_depths=cfg.depth_loss, |
| 451 | ) |
| 452 | self.valset = Dataset(self.parser, split="val") |
| 453 | self.scene_scale = self.parser.scene_scale * 1.1 * cfg.global_scale |
| 454 | print("Scene scale:", self.scene_scale) |
| 455 |