diff --git a/configs/graph_default.yaml b/configs/graph_default.yaml index 3f86c327e..361209285 100644 --- a/configs/graph_default.yaml +++ b/configs/graph_default.yaml @@ -86,12 +86,13 @@ graph: loss_coefficients: rgb_loss_coarse: 1.0 rgb_loss_fine: 1.0 - param_groups: - fields: - optimizer: - _target_: torch.optim.RAdam - lr: 0.0005 - scheduler: - _target_: pyrad.optimizers.optimizers.ExponentialDecaySchedule - lr_final: 0.000005 - max_steps: 1000000 + +optimizers: + fields: + optimizer: + _target_: torch.optim.RAdam + lr: 0.0005 + scheduler: + _target_: pyrad.optimizers.optimizers.ExponentialDecaySchedule + lr_final: 0.000005 + max_steps: 1000000 diff --git a/configs/graph_instant_ngp.yaml b/configs/graph_instant_ngp.yaml index 1b55b5ba6..5dc343f95 100644 --- a/configs/graph_instant_ngp.yaml +++ b/configs/graph_instant_ngp.yaml @@ -18,9 +18,9 @@ graph: loss_coefficients: # loss coefficients rgb_loss: 1.0 - # optimizer options for network - param_groups: - fields: - optimizer: - lr: 1e-2 - scheduler: null +# optimizer options for network +optimizers: + fields: + optimizer: + lr: 1e-2 + scheduler: null diff --git a/pyrad/data/dataloader.py b/pyrad/data/dataloader.py index c4e8f1f66..1577d3d8b 100644 --- a/pyrad/data/dataloader.py +++ b/pyrad/data/dataloader.py @@ -18,7 +18,7 @@ Data loader. import random from typing import Dict, List, Tuple, Union -from omegaconf import ListConfig +from omegaconf import DictConfig, ListConfig from torchtyping import TensorType @@ -27,7 +27,57 @@ from pyrad.cameras.rays import RayBundle from pyrad.data.image_dataset import ImageDataset from pyrad.data.image_sampler import ImageSampler from pyrad.data.pixel_sampler import PixelSampler -from pyrad.utils.misc import get_dict_to_torch +from pyrad.data.structs import DatasetInputs +from pyrad.data.utils import get_dataset_inputs_from_dataset_config +from pyrad.utils import profiler +from pyrad.utils.misc import get_dict_to_torch, instantiate_from_dict_config + + +@profiler.time_function +def setup_dataset_train(config: DictConfig, device: str) -> Tuple[DatasetInputs, "TrainDataloader"]: + """Helper method to load train dataset + Args: + config (DictConfig): Configuration of training dataset. + device (str): device to load the dataset to + + Returns: + Tuple[DatasetInputs, "TrainDataloader"]: returns both the dataset input information and associated dataloader + """ + dataset_inputs_train = get_dataset_inputs_from_dataset_config(**config.dataset_inputs_train, split="train") + # ImageDataset + image_dataset_train = instantiate_from_dict_config(config.image_dataset_train, **dataset_inputs_train.as_dict()) + # ImageSampler + image_sampler_train = instantiate_from_dict_config( + config.dataloader_train.image_sampler, dataset=image_dataset_train, device=device + ) + # PixelSampler + pixel_sampler_train = instantiate_from_dict_config(config.dataloader_train.pixel_sampler) + # Dataloader + dataloader_train = TrainDataloader(image_sampler_train, pixel_sampler_train) + return dataset_inputs_train, dataloader_train + + +@profiler.time_function +def setup_dataset_eval(config: DictConfig, test_mode: bool, device: str) -> Tuple[DatasetInputs, "EvalDataloader"]: + """Helper method to load test or val dataset based on test/train mode + Args: + config (DictConfig): Configuration of training dataset. + test_mode (bool): specifies whether you are training/testing mode, to load validation/test data + device (str): device to load the dataset to + + Returns: + Tuple[DatasetInputs, "TrainDataloader"]: returns both the dataset input information and associated dataloader + """ + eval_split = "test" if test_mode else "val" + dataset_inputs_eval = get_dataset_inputs_from_dataset_config(**config.dataset_inputs_eval, split=eval_split) + image_dataset_eval = instantiate_from_dict_config(config.image_dataset_eval, **dataset_inputs_eval.as_dict()) + dataloader_eval = instantiate_from_dict_config( + config.dataloader_eval, + image_dataset=image_dataset_eval, + device=device, + **dataset_inputs_eval.as_dict(), + ) + return dataset_inputs_eval, dataloader_eval class TrainDataloader: @@ -88,7 +138,7 @@ class FixedIndicesEvalDataloader(EvalDataloader): num_rays_per_chunk: int, image_indices: Union[List[int], ListConfig], device="cpu", - **kwargs + **kwargs, ): super().__init__(image_dataset, intrinsics, camera_to_world, num_rays_per_chunk, device, **kwargs) self.image_indices = image_indices diff --git a/pyrad/engine/trainer.py b/pyrad/engine/trainer.py index ee8543ce5..1bf4832e5 100644 --- a/pyrad/engine/trainer.py +++ b/pyrad/engine/trainer.py @@ -26,13 +26,12 @@ from torch.nn.parallel import DistributedDataParallel as DDP from torchtyping import TensorType from pyrad.cameras.rays import RayBundle -from pyrad.data.dataloader import EvalDataloader, TrainDataloader -from pyrad.data.utils import DatasetInputs, get_dataset_inputs_from_dataset_config -from pyrad.optimizers.optimizers import Optimizers +from pyrad.data.dataloader import EvalDataloader, setup_dataset_eval, setup_dataset_train +from pyrad.graphs.base import setup_graph +from pyrad.optimizers.optimizers import setup_optimizers from pyrad.utils import profiler, writer from pyrad.utils.callbacks import update_occupancy from pyrad.utils.decorators import check_main_thread -from pyrad.utils.misc import instantiate_from_dict_config from pyrad.utils.writer import EventName, TimeWriter logging.getLogger("PIL").setLevel(logging.WARNING) @@ -52,6 +51,7 @@ class Trainer: self.config = config self.local_rank = local_rank self.world_size = world_size + self.device = "cpu" if world_size == 0 else f"cuda:{local_rank}" # dataset variables self.dataloader_train = None self.dataloader_eval = None @@ -62,111 +62,21 @@ class Trainer: # logging variables writer.setup_event_writers(config) profiler.setup_profiler(config.logging) - self.device = "cpu" if self.world_size == 0 else f"cuda:{self.local_rank}" def setup(self, test_mode=False): """Setup the Trainer by calling other setup functions.""" - dataset_inputs_train = get_dataset_inputs_from_dataset_config( - **self.config.data.dataset_inputs_train, split="train" - ) - eval_split = "test" if test_mode else "val" - dataset_inputs_eval = get_dataset_inputs_from_dataset_config( - **self.config.data.dataset_inputs_eval, split=eval_split - ) - self.setup_dataset_train(dataset_inputs_train) - self.setup_dataset_eval(dataset_inputs_eval) - self.setup_graph(dataset_inputs_train) - - @profiler.time_function - def setup_dataset_train(self, dataset_inputs: DatasetInputs): - """_summary_""" - # ImageDataset - image_dataset_train = instantiate_from_dict_config( - self.config.data.image_dataset_train, **dataset_inputs.as_dict() - ) - # ImageSampler - image_sampler_train = instantiate_from_dict_config( - self.config.data.dataloader_train.image_sampler, dataset=image_dataset_train, device=self.device - ) - # PixelSampler - pixel_sampler_train = instantiate_from_dict_config(self.config.data.dataloader_train.pixel_sampler) - # Dataloader - self.dataloader_train = TrainDataloader(image_sampler_train, pixel_sampler_train) - - @profiler.time_function - def setup_dataset_eval(self, dataset_inputs: DatasetInputs): - """Helper method to load test or val dataset based on test/train mode""" - image_dataset_eval = instantiate_from_dict_config( - self.config.data.image_dataset_eval, **dataset_inputs.as_dict() - ) - self.dataloader_eval = instantiate_from_dict_config( - self.config.data.dataloader_eval, - image_dataset=image_dataset_eval, - device=self.device, - **dataset_inputs.as_dict(), - ) - - @profiler.time_function - def setup_graph(self, dataset_inputs: DatasetInputs): - """Setup the graph. The dataset inputs should be set with the training data. - - Args: - dataset_inputs (DatasetInputs): The inputs which will be used to define the camera parameters. - """ - self.graph = instantiate_from_dict_config(self.config.graph.network, **dataset_inputs.as_dict()) - self.graph.to(self.device) - - self.setup_optimizers() # NOTE(ethan): can this be before DDP? + dataset_inputs_train, self.dataloader_train = setup_dataset_train(self.config.data, device=self.device) + _, self.dataloader_eval = setup_dataset_eval(self.config.data, test_mode=test_mode, device=self.device) + self.graph = setup_graph(self.config.graph, dataset_inputs_train, device=self.device) + self.optimizers = setup_optimizers(self.config.optimizers, self.graph.get_param_groups()) if self.config.graph.resume_train.load_dir: - self.load_checkpoint(self.config.graph.resume_train) + self._load_checkpoint() if self.world_size > 1: self.graph = DDP(self.graph, device_ids=[self.local_rank]) dist.barrier(device_ids=[self.local_rank]) - def setup_optimizers(self): - """_summary_""" - self.optimizers = Optimizers(self.config.graph.param_groups, self.graph.get_param_groups()) - - def load_checkpoint(self, load_config: DictConfig) -> int: - """Load the checkpoint from the given path - - Args: - load_config (DictConfig): Configuration for loading a model. - - Returns: - int: step iteration associated with the loaded checkpoint - """ - load_path = os.path.join(load_config.load_dir, f"step-{load_config.load_step:09d}.ckpt") - assert os.path.exists(load_path), f"Checkpoint {load_path} does not exist" - loaded_state = torch.load(load_path, map_location="cpu") - self.graph.load_state_dict({key.replace("module.", ""): value for key, value in loaded_state["model"].items()}) - for k, v in loaded_state["optimizers"].items(): - self.optimizers.optimizers[k].load_state_dict(v) - self.start_step = loaded_state["step"] + 1 - logging.info("done loading checkpoint from %s", load_path) - - @check_main_thread - def save_checkpoint(self, output_dir: str, step: int) -> None: - """Save the model and optimizers - - Args: - output_dir (str): directory to save the checkpoint - step (int): number of steps in training for given checkpoint - """ - if not os.path.exists(output_dir): - os.makedirs(output_dir) - ckpt_path = os.path.join(output_dir, f"step-{step:09d}.ckpt") - torch.save( - { - "step": step, - "model": self.graph.module.state_dict() if hasattr(self.graph, "module") else self.graph.state_dict(), - "optimizers": {k: v.state_dict() for (k, v) in self.optimizers.optimizers.items()}, - }, - ckpt_path, - ) - @classmethod def get_aggregated_loss(cls, loss_dict: Dict[str, torch.tensor]): """Returns the aggregated losses and the scalar for calling .backwards() on. @@ -194,7 +104,7 @@ class Trainer: if step != 0 and step % self.config.logging.steps_per_log == 0: writer.put_dict(name="Loss/train-loss_dict", scalar_dict=loss_dict, step=step) if step != 0 and self.config.graph.steps_per_save and step % self.config.graph.steps_per_save == 0: - self.save_checkpoint(self.config.graph.model_dir, step) + self._save_checkpoint(self.config.graph.model_dir, step) if step % self.config.graph.steps_per_test == 0: self.eval_with_dataloader(self.dataloader_eval, step=step) self._write_out_storage(step) @@ -215,6 +125,38 @@ class Trainer: ): writer.write_out_storage() + def _load_checkpoint(self) -> None: + """helper function to load graph and optimizer from prespecified checkpoint""" + load_config = self.config.resume_train + load_path = os.path.join(load_config.load_dir, f"step-{load_config.load_step:09d}.ckpt") + assert os.path.exists(load_path), f"Checkpoint {load_path} does not exist" + loaded_state = torch.load(load_path, map_location="cpu") + self.start_step = loaded_state["step"] + 1 + # load the checkpoints for graph and optimizer + self.graph.load_graph(loaded_state) + self.optimizers.load_optimizers(loaded_state) + logging.info("done loading checkpoint from %s", load_path) + + @check_main_thread + def _save_checkpoint(self, output_dir: str, step: int) -> None: + """Save the model and optimizers + + Args: + output_dir (str): directory to save the checkpoint + step (int): number of steps in training for given checkpoint + """ + if not os.path.exists(output_dir): + os.makedirs(output_dir) + ckpt_path = os.path.join(output_dir, f"step-{step:09d}.ckpt") + torch.save( + { + "step": step, + "model": self.graph.module.state_dict() if hasattr(self.graph, "module") else self.graph.state_dict(), + "optimizers": {k: v.state_dict() for (k, v) in self.optimizers.optimizers.items()}, + }, + ckpt_path, + ) + @profiler.time_function def train_iteration( self, ray_indices: TensorType["num_rays", 3], batch: dict, step: int, _callback: List[Callable] = None diff --git a/pyrad/graphs/base.py b/pyrad/graphs/base.py index 61d456b27..b527f9267 100644 --- a/pyrad/graphs/base.py +++ b/pyrad/graphs/base.py @@ -17,7 +17,7 @@ The Graph module contains all trainable parameters. """ from abc import abstractmethod from collections import defaultdict -from typing import Dict, List, Union +from typing import Any, Dict, List, Union import torch from omegaconf import DictConfig @@ -27,11 +27,24 @@ from torchtyping import TensorType from pyrad.cameras.cameras import Camera from pyrad.cameras.rays import RayBundle -from pyrad.data.structs import SceneBounds +from pyrad.data.structs import DatasetInputs, SceneBounds from pyrad.graphs.modules.ray_generator import RayGenerator +from pyrad.utils import profiler from pyrad.utils.misc import get_masked_dict, instantiate_from_dict_config, is_not_none +@profiler.time_function +def setup_graph(config: DictConfig, dataset_inputs: DatasetInputs, device: str) -> "Graph": + """Setup the graph. The dataset inputs should be set with the training data. + + Args: + dataset_inputs (DatasetInputs): The inputs which will be used to define the camera parameters. + """ + graph = instantiate_from_dict_config(config.network, **dataset_inputs.as_dict()) + graph.to(device) + return graph + + class AbstractGraph(nn.Module): """Highest level graph class. Somewhat useful to lift code up and out of the way.""" @@ -178,3 +191,7 @@ class Graph(AbstractGraph): @abstractmethod def log_test_image_outputs(self, image_idx, step, batch, outputs): """Writes the test image outputs.""" + + def load_graph(self, loaded_state: Dict[str, Any]) -> None: + """Load the checkpoint from the given path""" + self.load_state_dict({key.replace("module.", ""): value for key, value in loaded_state["model"].items()}) diff --git a/pyrad/optimizers/optimizers.py b/pyrad/optimizers/optimizers.py index 38ef8e0c4..808f4da38 100644 --- a/pyrad/optimizers/optimizers.py +++ b/pyrad/optimizers/optimizers.py @@ -16,7 +16,7 @@ Optimizers class. """ -from typing import Dict, List +from typing import Any, Dict, List import numpy as np from hydra.utils import instantiate @@ -27,6 +27,19 @@ from torch.optim.lr_scheduler import LambdaLR from pyrad.utils import writer +def setup_optimizers(config: DictConfig, param_groups: Dict[str, List[Parameter]]) -> "Optimizers": + """Helper to set up the optimizers + + Args: + config (DictConfig): _description_ + param_groups (Dict[str, List[Parameter]]): _description_ + + Returns: + Optimizers: _description_ + """ + return Optimizers(config, param_groups) + + class ExponentialDecaySchedule(LambdaLR): """Exponential learning rate decay function. See https://github.com/google-research/google-research/blob/ @@ -108,3 +121,8 @@ class Optimizers: # TODO(ethan): clean this up. why is there indexing into a list? lr = scheduler.get_last_lr()[0] writer.put_scalar(name=f"learning_rate/{param_group_name}", scalar=lr, step=step) + + def load_optimizers(self, loaded_state: Dict[str, Any]) -> None: + """Helper to load the optimizer state from previous checkpoint""" + for k, v in loaded_state["optimizers"].items(): + self.optimizers[k].load_state_dict(v) diff --git a/pyrad/utils/writer.py b/pyrad/utils/writer.py index ea494dc20..4bef38e16 100644 --- a/pyrad/utils/writer.py +++ b/pyrad/utils/writer.py @@ -195,10 +195,11 @@ class Writer: class TimeWriter: """Timer context manager that calculates duration around wrapped functions""" - def __init__(self, writer, name, step=None): + def __init__(self, writer, name, step=None, write=True): self.writer = writer self.name = name self.step = step + self.write = write self.start = None self.duration = None @@ -210,13 +211,14 @@ class TimeWriter: def __exit__(self, *args): self.duration = time() - self.start update_step = self.step is not None - self.writer.put_time( - name=self.name, - duration=self.duration, - step=self.step if update_step else GLOBAL_BUFFER["max_iter"], - avg_over_steps=update_step, - update_eta=self.name == EventName.ITER_TRAIN_TIME, - ) + if self.write: + self.writer.put_time( + name=self.name, + duration=self.duration, + step=self.step if update_step else GLOBAL_BUFFER["max_iter"], + avg_over_steps=update_step, + update_eta=self.name == EventName.ITER_TRAIN_TIME, + ) @decorate_all([check_main_thread]) diff --git a/scripts/benchmarking/run_benchmark.py b/scripts/benchmarking/run_benchmark.py index 80f606f62..aad09fcab 100644 --- a/scripts/benchmarking/run_benchmark.py +++ b/scripts/benchmarking/run_benchmark.py @@ -4,15 +4,13 @@ run_benchmark.py import json import logging import os -from datetime import date -import torch from hydra import compose, initialize from hydra.core.global_hydra import GlobalHydra from omegaconf import DictConfig from tqdm import tqdm -from pyrad.engine.trainer import Trainer +from scripts.run_eval import run_inference BENCH = { "method": "vanilla_nerf", @@ -22,8 +20,16 @@ BENCH = { } -def load_best_ckpt(hydra_dir: str, config: DictConfig) -> str: - """helper function to update config with latest checkpoint in specified hydra_dir""" +def _load_best_ckpt(hydra_dir: str, config: DictConfig) -> str: + """helper function to update config with latest checkpoint in specified hydra_dir + + Args: + hydra_dir (str): base directory for the specified run (where the model directories are stored) + config (DictConfig): configuration for the specified run. + + Returns: + str: path to the most recent checkpoint sorted by timestamps and step number + """ model_dir = os.path.join(hydra_dir, config.graph.model_dir) latest_ckpt = os.listdir(model_dir) latest_ckpt.sort() @@ -35,6 +41,22 @@ def load_best_ckpt(hydra_dir: str, config: DictConfig) -> str: return os.path.join(model_dir, latest_ckpt) +def _load_hydra_config(hydra_dir: str) -> DictConfig: + """helper function to load the specified hydra config from the specified directory + + Args: + hydra_dir (str): base directory for the specified run (where the flattened config.yaml is stored) + + Returns: + DictConfig: returns the loaded hydra dictionary config + """ + basename = os.listdir(hydra_dir)[0] + hydra_dir = f"{hydra_dir}/{basename}" + initialize(config_path=os.path.join("../../", hydra_dir, ".hydra/")) + config = compose("config.yaml") + return config + + def main(): """Main function.""" benchmarks = {} @@ -42,22 +64,19 @@ def main(): method = BENCH["method"] benchmark_date = BENCH["benchmark_date"] for dataset in tqdm(BENCH["object_list"]): + # set up trainer, config, and checkpoint loading hydra_dir = f"{hydra_base_dir}/blender_{dataset}_{benchmark_date}/{method}/" - basename = os.listdir(hydra_dir)[0] - hydra_dir = f"{hydra_dir}/{basename}" - initialize(config_path=os.path.join("../../", hydra_dir, ".hydra/")) - config = compose("config.yaml") - ckpt = load_best_ckpt(hydra_dir, config) - trainer = Trainer(config, local_rank=0, world_size=1) - trainer.setup(test_mode=True) - avg_psnr = 0 - for step, (camera_ray_bundle, batch) in enumerate(trainer.dataloader_eval): - with torch.no_grad(): - psnr = trainer.test_image(camera_ray_bundle, batch, step=step) - avg_psnr = (step * avg_psnr + psnr) / (step + 1) - benchmarks[dataset] = (avg_psnr, ckpt) + config = _load_hydra_config(hydra_dir) + ckpt = _load_best_ckpt(hydra_dir, config) + + # run evaluation + avg_psnr, avg_rays_per_sec = run_inference(config) + benchmarks[dataset] = {"avg psnr": avg_psnr, "avg rays/s": avg_rays_per_sec, "checkpoint": ckpt} + + # reset hydra config GlobalHydra.instance().clear() + # output benchmark statistics to a json file benchmark_info = {"bench": BENCH, "results": benchmarks} timestamp = BENCH["benchmark_date"] json_file = os.path.join(BENCH["hydra_base_dir"], f"{timestamp}.json") diff --git a/scripts/run_eval.py b/scripts/run_eval.py index 4dec9adb5..ad4c33c4b 100644 --- a/scripts/run_eval.py +++ b/scripts/run_eval.py @@ -1,87 +1,90 @@ """ run_eval.py """ -import json import logging import os -from datetime import date - +from typing import Tuple +import hydra import torch -from hydra import compose, initialize -from hydra.core.global_hydra import GlobalHydra from omegaconf import DictConfig from tqdm import tqdm -from pyrad.engine.trainer import Trainer -from pyrad.cameras.cameras import Camera - -BENCH = { - "method": "vanilla_nerf", - "hydra_base_dir": "outputs/", - "benchmark_date": "05-26-2022", - "object_list": ["mic", "ficus", "chair", "hotdog", "materials", "drums", "ship", "lego"], -} +from pyrad.data.dataloader import setup_dataset_eval, setup_dataset_train +from pyrad.graphs.base import Graph, setup_graph +from pyrad.utils.writer import TimeWriter -def load_best_ckpt(hydra_dir: str, config: DictConfig) -> str: - """helper function to update config with latest checkpoint in specified hydra_dir""" - model_dir = os.path.join(hydra_dir, config.graph.model_dir) - latest_ckpt = os.listdir(model_dir) - latest_ckpt.sort() - latest_ckpt = latest_ckpt[-1] - step = os.path.splitext(latest_ckpt)[0].split("-")[-1] - config.graph.resume_train.load_dir = model_dir - config.graph.resume_train.load_step = int(step) - config.logging.writer.TensorboardWriter.save_dir = "/tmp/" - return os.path.join(model_dir, latest_ckpt) +def _update_avg(prev_avg: float, new_val: float, step: int) -> float: + """helper to calculate the running average + + Args: + prev_avg (float): previous average value + new_val (float): new value to update the average with + step (int): current step number + + Returns: + float: new updated average + """ + return (step * prev_avg + new_val) / (step + 1) -def main(): +def _load_checkpoint(config: DictConfig, graph: Graph) -> None: + """Helper function to load checkpointed graph + + Args: + config (DictConfig): Configuration of graph to load + graph (Graph): Graph instance of which to load weights + """ + load_path = os.path.join(config.load_dir, f"step-{config.load_step:09d}.ckpt") + assert os.path.exists(load_path), f"Checkpoint {load_path} does not exist" + loaded_state = torch.load(load_path, map_location="cpu") + graph.load_checkpoint(loaded_state) + logging.info("done loading checkpoint from %s", load_path) + + +def run_inference(config: DictConfig, local_rank: int = 0, world_size: int = 1) -> Tuple[float, float]: + """helper function to run inference given config specifications (also used in benchmarking) + + Args: + config (DictConfig): Configuration for loading the evaluation. + + Returns: + Tuple[float, float]: returns both the avg psnr and avg rays per second + """ + device = "cpu" if world_size == 0 else f"cuda:{local_rank}" + # setup graph and dataset + dataset_inputs_train, _ = setup_dataset_train(config.data, device=device) + _, dataloader_eval = setup_dataset_eval(config.data, test_mode=True, device=device) + graph = setup_graph(config.graph, dataset_inputs_train, device=device) + + # load checkpointed information + _load_checkpoint(config.graph.resume_train, graph) + + # calculate average psnr across test dataset + # TODO(ethan): trajector specification + avg_psnr = 0 + avg_rays_per_sec = 0 + for step, (camera_ray_bundle, batch) in tqdm(enumerate(dataloader_eval)): + with TimeWriter(writer=None, name=None, write=False) as t: + with torch.no_grad(): + image_idx = int(camera_ray_bundle.camera_indices[0, 0]) + outputs = graph.get_outputs_for_camera_ray_bundle(camera_ray_bundle) + psnr = graph.log_test_image_outputs(image_idx, step, batch, outputs) + avg_rays_per_sec = _update_avg(avg_rays_per_sec, camera_ray_bundle.origins.shape[0] / t.duration, step) + avg_psnr = _update_avg(avg_psnr, psnr, step) + return avg_psnr, avg_rays_per_sec + + +@hydra.main(config_path="../configs", config_name="graph_default.yaml") +def main(config: DictConfig): """Main function.""" - # benchmarks = {} - # hydra_base_dir = BENCH["hydra_base_dir"] - # method = BENCH["method"] - # benchmark_date = BENCH["benchmark_date"] - # for dataset in tqdm(BENCH["object_list"]): + assert config.graph.resume_train.load_dir, "Please specify checkpoint load path" + assert config.graph.resume_train.load_step, "Please specify checkpoint step to load" - # hydra_dir = f"{hydra_base_dir}/blender_{dataset}_{benchmark_date}/{method}/" - # basename = os.listdir(hydra_dir)[0] - # hydra_dir = f"{hydra_dir}/{basename}" - # initialize(config_path=os.path.join("../../", hydra_dir, ".hydra/")) - # config = compose("configs/flat.yaml") + avg_psnr, avg_rays_per_sec = run_inference(config) - with initialize(version_base=None, config_path="../configs"): - config = compose(config_name="flat.yaml") - trainer = Trainer(config, local_rank=0, world_size=1) - trainer.setup(test_mode=False) - - # create a camera trajectory - print("Running eval.") - # camera = Camera() - # trainer.graph.eval() - # camera_ray_bundle = camera.get_camera_ray_bundle(device=trainer.device) - image_idx = 0 - camera_ray_bundle, batch = trainer.dataloader_eval.get_data_from_image_idx(image_idx) - image_height, image_width = camera_ray_bundle.origins.shape[:2] - print(image_height, image_width) - outputs = trainer.graph.get_outputs_for_camera_ray_bundle(camera_ray_bundle) - - # print(outputs) - - # avg_psnr = 0 - # for step, image_idx in enumerate(config.data.val_image_indices): - # with torch.no_grad(): - # psnr = trainer.test_image(image_idx=image_idx, step=step) - # avg_psnr = (step * avg_psnr + psnr) / (step + 1) - # benchmarks[dataset] = (avg_psnr, ckpt) - # GlobalHydra.instance().clear() - - # benchmark_info = {"bench": BENCH, "results": benchmarks} - # timestamp = date.today().strftime("%b-%d-%Y") - # json_file = os.path.join(BENCH["hydra_base_dir"], f"{timestamp}.json") - # with open(json_file, "w", encoding="utf8") as f: - # json.dump(benchmark_info, f, indent=2) - # logging.info("saved benchmark results to %s", json_file) + print(f"Avg. PSNR: {avg_psnr:0.4f}") + print(f"Avg. Rays per sec: {avg_rays_per_sec:0.4f}") if __name__ == "__main__": diff --git a/scripts/run_train.py b/scripts/run_train.py index c37b94766..791b1ccde 100644 --- a/scripts/run_train.py +++ b/scripts/run_train.py @@ -116,8 +116,8 @@ def _train(local_rank: int, world_size: int, config: DictConfig) -> Any: Args: local_rank (int): _description_ - world_size (int): _description_ - config (DictConfig): _description_ + local_rank (int): current rank of process + world_size (int): total number of gpus available Returns: Any: TODO(): determine the return type