Function main

examples/simple_trainer.py:1465–1498

(local_rank: int, world_rank, world_size: int, cfg: Config)

def main(local_rank: int, world_rank, world_size: int, cfg: Config):
    if world_size > 1 and not cfg.disable_viewer:
        cfg.disable_viewer = True
        if world_rank == 0:
            print("Viewer is disabled in distributed training.")

    runner = Runner(local_rank, world_rank, world_size, cfg)

    if cfg.mode == "train":
        runner.train()
    elif cfg.mode == "compress":
        if cfg.ckpt is None:
            raise ValueError("ckpt parameter is required for compress mode")
        # Load checkpoints
        ckpts = [
            torch.load(file, map_location=runner.device, weights_only=True)
            for file in cfg.ckpt
        ]
        for k in runner.splats.keys():
            runner.splats[k].data = torch.cat([ckpt["splats"][k] for ckpt in ckpts])
        step = ckpts[0]["step"]
        # save ply
        runner.save_params_into_ply_file()
        # eval and render traj. of uncompressed splats
        # runner.eval(step=step)
        # runner.render_traj(step=step)

        runner.run_post_training_compression(step=step)
    else:
        raise ValueError(f"Unknown mode: {cfg.mode}")

    if not cfg.disable_viewer:
        print("Viewer running... Ctrl+C to exit.")
        time.sleep(1000000)


if __name__ == "__main__":