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

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

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)

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)