From 91ef54963c43beb34f13301f8496faf2f0de8a2e Mon Sep 17 00:00:00 2001 From: evonneng Date: Wed, 15 Jun 2022 09:06:43 -0700 Subject: [PATCH] refactoring configs; pulled network out of graph; created trainer heirarchy (#51) * changed config graph structure * removed network udner graphs; added new heirarchy trainer * fixed callback in occupancy graph * lint * changed register callback to setter method --- README.md | 2 +- configs/graph_default.yaml | 41 ++++++++++--------- configs/graph_instant_ngp.yaml | 13 +++--- configs/graph_mipnerf.yaml | 7 ++-- configs/graph_nerfw.yaml | 21 +++++----- configs/graph_semantic_nerf.yaml | 19 ++++----- configs/tests/test_dryrun.yaml | 11 ++--- docs/tooling/logging.md | 4 +- docs/tutorials/creating_graphs.md | 23 +++++++---- pyrad/engine/trainer.py | 33 ++++++++------- pyrad/fields/base.py | 6 +-- .../fields/occupancy_fields/occupancy_grid.py | 5 +-- pyrad/graphs/base.py | 7 +++- pyrad/graphs/instant_ngp.py | 11 +++++ pyrad/utils/callbacks.py | 41 ++++++++++++------- pyrad/utils/writer.py | 28 ++++++------- scripts/benchmarking/launch_train_blender.sh | 6 +-- scripts/benchmarking/run_benchmark.py | 11 ++--- scripts/run_eval.py | 8 ++-- 19 files changed, 163 insertions(+), 134 deletions(-) diff --git a/README.md b/README.md index 33574e712..522e28dd5 100644 --- a/README.md +++ b/README.md @@ -78,7 +78,7 @@ python scripts/run_train.py data/dataset=blender_lego python scripts/run_train.py data/dataset=friends_TBBT-big_living_room # Run with different datasets and config changes -python scripts/run_train.py data/dataset=friends_TBBT-big_living_room graph.network.far_plane=14 +python scripts/run_train.py data/dataset=friends_TBBT-big_living_room graph.far_plane=14 # [Experimental] Speed up the dataloading pipeline by caching DatasetInputs. python scripts/run_data_preprocessor.py diff --git a/configs/graph_default.yaml b/configs/graph_default.yaml index 361209285..b87155e30 100644 --- a/configs/graph_default.yaml +++ b/configs/graph_default.yaml @@ -16,15 +16,24 @@ logging: # writer logs losses and images per iteration writer: TensorboardWriter: - save_dir: "./" + log_dir: "./" LocalWriter: - save_dir: "./" + log_dir: "./" stats_to_track: [ITER_LOAD_TIME, ITER_TRAIN_TIME, RAYS_PER_SEC, CURR_TEST_PSNR] max_log_size: 10 # if 0, logs everything with no erasing # profiler logs run times of functions and prints at end of training enable_profiler: True +trainer: + model_dir: /tmp/pyrad_models/ + steps_per_save: 100 + steps_per_test: 1000 + max_num_iterations: 1000000 + resume_train: + load_dir: # pretrained model directory + load_step: # pretrained model's step number + experiment_name: blender_lego method_name: vanilla_nerf @@ -68,24 +77,16 @@ data: # num_rays_per_chunk: 4096 graph: - model_dir: /tmp/pyrad_models/ - steps_per_save: 100 - steps_per_test: 1000 - max_num_iterations: 1000000 - resume_train: - load_dir: # pretrained model directory - load_step: # pretrained model's step number - network: - _target_: pyrad.graphs.vanilla_nerf.NeRFGraph - collider_config: - _target_: pyrad.graphs.modules.scene_colliders.NearFarCollider - near_plane: 2.0 - far_plane: 6.0 - num_coarse_samples: 64 - num_importance_samples: 128 - loss_coefficients: - rgb_loss_coarse: 1.0 - rgb_loss_fine: 1.0 + _target_: pyrad.graphs.vanilla_nerf.NeRFGraph + collider_config: + _target_: pyrad.graphs.modules.scene_colliders.NearFarCollider + near_plane: 2.0 + far_plane: 6.0 + num_coarse_samples: 64 + num_importance_samples: 128 + loss_coefficients: + rgb_loss_coarse: 1.0 + rgb_loss_fine: 1.0 optimizers: fields: diff --git a/configs/graph_instant_ngp.yaml b/configs/graph_instant_ngp.yaml index 5dc343f95..924c50c0a 100644 --- a/configs/graph_instant_ngp.yaml +++ b/configs/graph_instant_ngp.yaml @@ -10,13 +10,12 @@ data: num_rays_per_chunk: 16384 # for evaluation graph: - network: - _target_: pyrad.graphs.instant_ngp.NGPGraph - collider_config: - _target_: pyrad.graphs.modules.scene_colliders.AABBBoxCollider - field_implementation: tcnn # torch, tcnn, ... - loss_coefficients: # loss coefficients - rgb_loss: 1.0 + _target_: pyrad.graphs.instant_ngp.NGPGraph + collider_config: + _target_: pyrad.graphs.modules.scene_colliders.AABBBoxCollider + field_implementation: tcnn # torch, tcnn, ... + loss_coefficients: # loss coefficients + rgb_loss: 1.0 # optimizer options for network optimizers: diff --git a/configs/graph_mipnerf.yaml b/configs/graph_mipnerf.yaml index 2529a6f72..1317b735e 100644 --- a/configs/graph_mipnerf.yaml +++ b/configs/graph_mipnerf.yaml @@ -10,7 +10,6 @@ method_name: mipnerf # rgb = sigmoid(x) rgb = (1 + 2e) / (1 + exp(-x)) -e graph: - network: - _target_: pyrad.graphs.mipnerf.MipNerfGraph - num_coarse_samples: 128 - num_importance_samples: 128 \ No newline at end of file + _target_: pyrad.graphs.mipnerf.MipNerfGraph + num_coarse_samples: 128 + num_importance_samples: 128 \ No newline at end of file diff --git a/configs/graph_nerfw.yaml b/configs/graph_nerfw.yaml index cf15b5e67..22d9e0f4f 100644 --- a/configs/graph_nerfw.yaml +++ b/configs/graph_nerfw.yaml @@ -21,14 +21,13 @@ data: downscale_factor: 1 graph: - network: - _target_: pyrad.graphs.nerfw.NerfWGraph - collider_config: - _target_: pyrad.graphs.modules.scene_colliders.AABBBoxCollider - num_coarse_samples: 64 - num_importance_samples: 64 - loss_coefficients: - rgb_loss_coarse: 1.0 - rgb_loss_fine: 1.0 - uncertainty_loss: 1.0 - density_loss: 0.01 + _target_: pyrad.graphs.nerfw.NerfWGraph + collider_config: + _target_: pyrad.graphs.modules.scene_colliders.AABBBoxCollider + num_coarse_samples: 64 + num_importance_samples: 64 + loss_coefficients: + rgb_loss_coarse: 1.0 + rgb_loss_fine: 1.0 + uncertainty_loss: 1.0 + density_loss: 0.01 diff --git a/configs/graph_semantic_nerf.yaml b/configs/graph_semantic_nerf.yaml index f41ce90a2..80e53a5e3 100644 --- a/configs/graph_semantic_nerf.yaml +++ b/configs/graph_semantic_nerf.yaml @@ -21,13 +21,12 @@ data: downscale_factor: 1 graph: - network: - _target_: pyrad.graphs.semantic_nerf.SemanticNerfGraph - collider_config: - _target_: pyrad.graphs.modules.scene_colliders.AABBBoxCollider - num_coarse_samples: 64 - num_importance_samples: 64 - loss_coefficients: - rgb_loss_coarse: 1.0 - rgb_loss_fine: 1.0 - semantic_loss_fine: 0.05 + _target_: pyrad.graphs.semantic_nerf.SemanticNerfGraph + collider_config: + _target_: pyrad.graphs.modules.scene_colliders.AABBBoxCollider + num_coarse_samples: 64 + num_importance_samples: 64 + loss_coefficients: + rgb_loss_coarse: 1.0 + rgb_loss_fine: 1.0 + semantic_loss_fine: 0.05 diff --git a/configs/tests/test_dryrun.yaml b/configs/tests/test_dryrun.yaml index aafb23c16..49b7c62fc 100644 --- a/configs/tests/test_dryrun.yaml +++ b/configs/tests/test_dryrun.yaml @@ -4,6 +4,9 @@ defaults: machine_config: num_gpus: 0 +trainer: + max_num_iterations: 2 + data: dataset_inputs_train: data_directory: tests/data/lego_test @@ -20,15 +23,13 @@ data: downscale_factor: 16 graph: - max_num_iterations: 2 - network: - num_coarse_samples: 4 - num_importance_samples: 4 + num_coarse_samples: 4 + num_importance_samples: 4 logging: writer: TensorboardWriter: - save_dir: "/tmp/" + log_dir: "/tmp/" hydra: searchpath: diff --git a/docs/tooling/logging.md b/docs/tooling/logging.md index 25bf6979e..f3514f171 100644 --- a/docs/tooling/logging.md +++ b/docs/tooling/logging.md @@ -11,9 +11,9 @@ logging: max_buffer_size: 20 # defines how many steps to average over (if reporting averages) writer: # definition of all writers you want to use TensorboardWriter: - save_dir: "./" # must specify the output path for all writers + log_dir: "./" # must specify the output path for all writers LocalWriter: - save_dir: "./" + log_dir: "./" stats_to_track: [ITER_LOAD_TIME, ITER_TRAIN_TIME, RAYS_PER_SEC, CURR_TEST_PSNR] max_log_size: 10 diff --git a/docs/tutorials/creating_graphs.md b/docs/tutorials/creating_graphs.md index 5175ed69e..ad0ae128c 100644 --- a/docs/tutorials/creating_graphs.md +++ b/docs/tutorials/creating_graphs.md @@ -115,20 +115,27 @@ To set up the graph config, create a new yaml under `configs/`. # configs/vanilla_nerf.yaml defaults: - - default_setup # inherit the basic yaml heirarchy + - graph_default # inherit the basic yaml heirarchy - _self_ +logging: + # + +trainer: + # + experiment_name: blender_lego method_name: vanilla_nerf -graph: - network: - _target_: pyrad.graph.vanilla_nerf.NeRFGraph # set the target to the graph you defined - - # - data: - # + # + +graph: + _target_: pyrad.graph.vanilla_nerf.NeRFGraph # set the target to the graph you defined + # + +optimizers: + # ``` ## Training and testing diff --git a/pyrad/engine/trainer.py b/pyrad/engine/trainer.py index 1bf4832e5..f13e1c2ed 100644 --- a/pyrad/engine/trainer.py +++ b/pyrad/engine/trainer.py @@ -17,7 +17,7 @@ Code to train model. """ import logging import os -from typing import Callable, Dict, List +from typing import Dict import torch import torch.distributed as dist @@ -30,7 +30,6 @@ from pyrad.data.dataloader import EvalDataloader, setup_dataset_eval, setup_data 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.writer import EventName, TimeWriter @@ -70,13 +69,15 @@ class Trainer: 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: + if self.config.trainer.resume_train.load_dir: 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]) + self.graph.register_callbacks() + @classmethod def get_aggregated_loss(cls, loss_dict: Dict[str, torch.tensor]): """Returns the aggregated losses and the scalar for calling .backwards() on. @@ -91,21 +92,21 @@ class Trainer: def train(self) -> None: """_summary_""" with TimeWriter(writer, EventName.TOTAL_TRAIN_TIME): - num_iterations = self.config.graph.max_num_iterations + num_iterations = self.config.trainer.max_num_iterations iter_dataloader_train = iter(self.dataloader_train) for step in range(self.start_step, self.start_step + num_iterations): with TimeWriter(writer, EventName.ITER_LOAD_TIME, step=step): ray_indices, batch = next(iter_dataloader_train) with TimeWriter(writer, EventName.ITER_TRAIN_TIME, step=step) as t: - loss_dict = self.train_iteration(ray_indices, batch, step, _callback=[update_occupancy]) + loss_dict = self.train_iteration(ray_indices, batch, step) writer.put_scalar(name=EventName.RAYS_PER_SEC, scalar=ray_indices.shape[0] / t.duration, step=step) 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) - if step % self.config.graph.steps_per_test == 0: + if step != 0 and self.config.trainer.steps_per_save and step % self.config.trainer.steps_per_save == 0: + self._save_checkpoint(self.config.trainer.model_dir, step) + if step % self.config.trainer.steps_per_test == 0: self.eval_with_dataloader(self.dataloader_eval, step=step) self._write_out_storage(step) @@ -119,9 +120,9 @@ class Trainer: """ if ( step % self.config.logging.steps_per_log == 0 - or (self.config.graph.steps_per_save and step % self.config.graph.steps_per_save == 0) - or step % self.config.graph.steps_per_test == 0 - or step == self.config.graph.max_num_iterations + or (self.config.trainer.steps_per_save and step % self.config.trainer.steps_per_save == 0) + or step % self.config.trainer.steps_per_test == 0 + or step == self.config.trainer.max_num_iterations ): writer.write_out_storage() @@ -158,9 +159,7 @@ class Trainer: ) @profiler.time_function - def train_iteration( - self, ray_indices: TensorType["num_rays", 3], batch: dict, step: int, _callback: List[Callable] = None - ) -> Dict[str, float]: + def train_iteration(self, ray_indices: TensorType["num_rays", 3], batch: dict, step: int) -> Dict[str, float]: """Run one iteration with a batch of inputs. Args: @@ -177,9 +176,9 @@ class Trainer: loss.backward() self.optimizers.optimizer_step_all() self.optimizers.scheduler_step_all(step) - if _callback: - for _func in _callback: - _func(self.graph) + if self.graph.callbacks: + for _func in self.graph.callbacks: + _func.after_step(step) return loss_dict @profiler.time_function diff --git a/pyrad/fields/base.py b/pyrad/fields/base.py index b736d48ce..7c892839a 100644 --- a/pyrad/fields/base.py +++ b/pyrad/fields/base.py @@ -38,9 +38,9 @@ class Field(nn.Module): frustums=Frustums( origins=positions, directions=torch.ones_like(positions), - frustum_starts=torch.zeros((*positions.shape[:-1], 1)), - frustum_ends=torch.zeros((*positions.shape[:-1], 1)), - pixel_area=torch.ones((*positions.shape[:-1], 1)), + frustum_starts=torch.zeros_like((positions[..., :1])), + frustum_ends=torch.zeros_like((positions[..., :1])), + pixel_area=torch.ones_like((positions[..., :1])), ) ) density, _ = self.get_density(point_samples) diff --git a/pyrad/fields/occupancy_fields/occupancy_grid.py b/pyrad/fields/occupancy_fields/occupancy_grid.py index c900b7e97..73b858ab0 100644 --- a/pyrad/fields/occupancy_fields/occupancy_grid.py +++ b/pyrad/fields/occupancy_fields/occupancy_grid.py @@ -28,7 +28,7 @@ class OccupancyGrid(nn.Module): """Module to keep track of the density and occupancy.""" def __init__( - self, num_cascades: int = 1, resolution: int = 128, aabb=None, density_fn=None, update_every_num_iters: int = 16 + self, num_cascades: int = 1, resolution: int = 128, aabb=None, update_every_num_iters: int = 16 ) -> None: super().__init__() assert is_not_none(aabb), "The axis-aligned bounding box aabb is not defined!" @@ -42,9 +42,7 @@ class OccupancyGrid(nn.Module): self.register_buffer("occupancy_grid", occupancy_grid) self.mean_density = 0.0 - self.density_fn = density_fn self.update_every_num_iters = update_every_num_iters - self.iteration_count = 0 def reset(self): """Zero out the occupancy grid.""" @@ -67,7 +65,6 @@ class OccupancyGrid(nn.Module): """ # TODO(ethan): remove this pylint disable # pylint: disable=invalid-name - device = self.occupancy_grid.device assert self.num_cascades == 1, "Still need to implement code for multiple cascades." diff --git a/pyrad/graphs/base.py b/pyrad/graphs/base.py index b527f9267..458e8fc1e 100644 --- a/pyrad/graphs/base.py +++ b/pyrad/graphs/base.py @@ -40,7 +40,7 @@ def setup_graph(config: DictConfig, dataset_inputs: DatasetInputs, device: str) 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 = instantiate_from_dict_config(config, **dataset_inputs.as_dict()) graph.to(device) return graph @@ -99,6 +99,11 @@ class Graph(AbstractGraph): self.populate_collider() self.populate_fields() self.populate_misc_modules() # populate the modules + self.callbacks = None + + def register_callbacks(self): # pylint:disable=no-self-use + """Option to register callback for training functions""" + self.callbacks = [] def populate_collider(self): """Set the scene bounds collider to use.""" diff --git a/pyrad/graphs/instant_ngp.py b/pyrad/graphs/instant_ngp.py index f2fe9b056..ac8283b21 100644 --- a/pyrad/graphs/instant_ngp.py +++ b/pyrad/graphs/instant_ngp.py @@ -33,6 +33,7 @@ from pyrad.graphs.modules.ray_sampler import PDFSampler, UniformSampler from pyrad.utils import colors from pyrad.cameras.rays import RayBundle from pyrad.utils import visualization, writer +from pyrad.utils.callbacks import Callback class NGPGraph(Graph): @@ -45,6 +46,16 @@ class NGPGraph(Graph): self.field = None super().__init__(intrinsics=intrinsics, camera_to_world=camera_to_world, **kwargs) + def register_callbacks(self) -> None: + """defining callbacks to run after every training iteration""" + self.callbacks = [ + Callback( + self.occupancy_grid.update_every_num_iters, + self.occupancy_grid.update_occupancy_grid, + density_fn=self.field.density_fn, + ) + ] + def populate_fields(self): """Set the fields.""" # torch or tiny-cuda-nn version diff --git a/pyrad/utils/callbacks.py b/pyrad/utils/callbacks.py index f82ab6992..691c4e0a5 100644 --- a/pyrad/utils/callbacks.py +++ b/pyrad/utils/callbacks.py @@ -16,20 +16,31 @@ Callback functions for training iterations """ -from pyrad.graphs.base import Graph -from pyrad.fields.occupancy_fields.occupancy_grid import OccupancyGrid +from typing import Callable -def update_occupancy(graph: Graph): - """Update the occupancy grid depending on iteration count after every train step""" - if not hasattr(graph, "occupancy_pointers"): - occupancy_grids = [] - for _, module in graph.__dict__["_modules"].items(): - if isinstance(module, OccupancyGrid): - occupancy_grids.append(module) - graph.occupancy_pointers = occupancy_grids - for grid in graph.occupancy_pointers: - grid.iteration_count += 1 - if grid.iteration_count % grid.update_every_num_iters == 0: - if grid.density_fn is not None: - grid.update_occupancy_grid(grid.density_fn) +class Callback: + """Callback Object""" + + def __init__(self, update_every_num_iters: int, func: Callable, *args, **kwargs): + self.update_every_num_iters = update_every_num_iters + self.func = func + self.args = args + self.kwargs = kwargs + + def before_step(self, step: int): + """callback to run before training step + + Args: + step (int): current iteration step + """ + raise NotImplementedError + + def after_step(self, step: int): + """callback to run after training step + + Args: + step (int): current iteration step + """ + if step % self.update_every_num_iters == 0: + self.func(*self.args, **self.kwargs) diff --git a/pyrad/utils/writer.py b/pyrad/utils/writer.py index 4bef38e16..8143c4688 100644 --- a/pyrad/utils/writer.py +++ b/pyrad/utils/writer.py @@ -141,13 +141,13 @@ def setup_event_writers(config: DictConfig) -> None: writer_class = getattr(sys.modules[__name__], writer_type) writer_config = logging_configs[writer_type] if writer_type == "LocalWriter": - curr_writer = writer_class(writer_config.save_dir, writer_config.stats_to_track, writer_config.max_log_size) + curr_writer = writer_class(writer_config.log_dir, writer_config.stats_to_track, writer_config.max_log_size) else: - curr_writer = writer_class(writer_config.save_dir) + curr_writer = writer_class(writer_config.log_dir) EVENT_WRITERS.append(curr_writer) ## configure all the global buffer basic information - GLOBAL_BUFFER["max_iter"] = config.graph.max_num_iterations + GLOBAL_BUFFER["max_iter"] = config.trainer.max_num_iterations GLOBAL_BUFFER["max_buffer_size"] = config.logging.max_buffer_size GLOBAL_BUFFER["steps_per_log"] = config.logging.steps_per_log GLOBAL_BUFFER["events"] = {} @@ -156,8 +156,8 @@ def setup_event_writers(config: DictConfig) -> None: class Writer: """Writer class""" - def __init__(self, save_dir: str): - self.save_dir = save_dir + def __init__(self, log_dir: str): + self.log_dir = log_dir @abstractmethod def write_image(self, name: str, image: TensorType["H", "W", "C"], step: int) -> None: @@ -225,9 +225,9 @@ class TimeWriter: class WandbWriter(Writer): """WandDB Writer Class""" - def __init__(self, save_dir: str): - super().__init__(save_dir) - wandb.init(dir=save_dir) + def __init__(self, log_dir: str): + super().__init__(log_dir) + wandb.init(dir=log_dir) def write_image(self, name: str, image: TensorType["H", "W", "C"], step: int) -> None: """_summary_ @@ -253,9 +253,9 @@ class WandbWriter(Writer): class TensorboardWriter(Writer): """Tensorboard Writer Class""" - def __init__(self, save_dir: str): - super().__init__(save_dir) - self.tb_writer = SummaryWriter(log_dir=self.save_dir) + def __init__(self, log_dir: str): + super().__init__(log_dir) + self.tb_writer = SummaryWriter(log_dir=self.log_dir) def write_image(self, name: str, image: TensorType["H", "W", "C"], step: int) -> None: """_summary_ @@ -310,13 +310,13 @@ def _format_time(seconds): class LocalWriter(Writer): """Local Writer Class""" - def __init__(self, save_dir: str, stats_to_track: ListConfig, max_log_size: int = 0): + def __init__(self, log_dir: str, stats_to_track: ListConfig, max_log_size: int = 0): """ Args: stats_to_track (ListConfig): the names of stats that should be logged. max_log size (int): max number of lines that will be logged to teminal. """ - super().__init__(save_dir) + super().__init__(log_dir) self.stats_to_track = [EventName[name].value for name in stats_to_track] self.max_log_size = max_log_size self.keys = set() @@ -325,7 +325,7 @@ class LocalWriter(Writer): def write_image(self, name: str, image: TensorType["H", "W", "C"], step: int) -> None: if name in self.stats_to_track: image = to8b(image) - image_path = os.path.join(self.save_dir, f"{name}.jpg") + image_path = os.path.join(self.log_dir, f"{name}.jpg") imageio.imwrite(image_path, np.uint8(image.cpu().numpy() * 255.0)) def write_scalar(self, name: str, scalar: float, step: int) -> None: diff --git a/scripts/benchmarking/launch_train_blender.sh b/scripts/benchmarking/launch_train_blender.sh index 0c1edf2ce..0739c5967 100755 --- a/scripts/benchmarking/launch_train_blender.sh +++ b/scripts/benchmarking/launch_train_blender.sh @@ -43,9 +43,9 @@ for dataset in ${DATASETS[@]}; do data.dataset_inputs_train.data_directory=data/blender/${dataset} \ data.dataset_inputs_eval.data_directory=data/blender/${dataset} \ experiment_name=blender_${dataset}_${tag} \ - graph.model_dir=pyrad_models/ \ - graph.steps_per_save=25000 \ - graph.max_num_iterations=2000000 \ + trainer.model_dir=pyrad_models/ \ + trainer.steps_per_save=25000 \ + trainer.max_num_iterations=2000000 \ logging.enable_profiler=False & echo "Launched ${METHOD} ${dataset} on gpu ${GPU_IDX[$idx]}, ${tag}" done \ No newline at end of file diff --git a/scripts/benchmarking/run_benchmark.py b/scripts/benchmarking/run_benchmark.py index aad09fcab..623a20bd2 100644 --- a/scripts/benchmarking/run_benchmark.py +++ b/scripts/benchmarking/run_benchmark.py @@ -30,14 +30,13 @@ def _load_best_ckpt(hydra_dir: str, config: DictConfig) -> str: 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) + model_dir = os.path.join(hydra_dir, config.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/" + config.resume_train.load_dir = model_dir + config.resume_train.load_step = int(step) return os.path.join(model_dir, latest_ckpt) @@ -54,6 +53,8 @@ def _load_hydra_config(hydra_dir: str) -> DictConfig: hydra_dir = f"{hydra_dir}/{basename}" initialize(config_path=os.path.join("../../", hydra_dir, ".hydra/")) config = compose("config.yaml") + ## set the tensorboard to a /tmp/ output file for now + config.logging.writer.TensorboardWriter.log_dir = "/tmp/" return config @@ -67,7 +68,7 @@ def main(): # set up trainer, config, and checkpoint loading hydra_dir = f"{hydra_base_dir}/blender_{dataset}_{benchmark_date}/{method}/" config = _load_hydra_config(hydra_dir) - ckpt = _load_best_ckpt(hydra_dir, config) + ckpt = _load_best_ckpt(hydra_dir, config.trainer) # run evaluation avg_psnr, avg_rays_per_sec = run_inference(config) diff --git a/scripts/run_eval.py b/scripts/run_eval.py index ad4c33c4b..8a84edbb5 100644 --- a/scripts/run_eval.py +++ b/scripts/run_eval.py @@ -38,7 +38,7 @@ def _load_checkpoint(config: DictConfig, graph: Graph) -> None: 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) + graph.load_graph(loaded_state) logging.info("done loading checkpoint from %s", load_path) @@ -58,7 +58,7 @@ def run_inference(config: DictConfig, local_rank: int = 0, world_size: int = 1) graph = setup_graph(config.graph, dataset_inputs_train, device=device) # load checkpointed information - _load_checkpoint(config.graph.resume_train, graph) + _load_checkpoint(config.trainer.resume_train, graph) # calculate average psnr across test dataset # TODO(ethan): trajector specification @@ -78,8 +78,8 @@ def run_inference(config: DictConfig, local_rank: int = 0, world_size: int = 1) @hydra.main(config_path="../configs", config_name="graph_default.yaml") def main(config: DictConfig): """Main function.""" - 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" + assert config.trainer.resume_train.load_dir, "Please specify checkpoint load path" + assert config.trainer.resume_train.load_step, "Please specify checkpoint step to load" avg_psnr, avg_rays_per_sec = run_inference(config)