From 1f1631cec6c4853fdbf1c7d8775b571973fd9fd5 Mon Sep 17 00:00:00 2001
From: evonneng <evonne_ng@berkeley.edu>
Date: Mon, 15 Aug 2022 21:31:43 +0000
Subject: [PATCH] removing callable modules from script/

---
 nerfactory/engine/trainer.py | 19 ++++++++++++++++++-
 scripts/__init__.py          |  0
 scripts/run_train.py         | 21 ++-------------------
 tests/test_run_train.py      | 20 ++++++++++++--------
 4 files changed, 32 insertions(+), 28 deletions(-)
 delete mode 100644 scripts/__init__.py

diff --git a/nerfactory/engine/trainer.py b/nerfactory/engine/trainer.py
index 181351f42..6795327ca 100644
--- a/nerfactory/engine/trainer.py
+++ b/nerfactory/engine/trainer.py
@@ -19,7 +19,7 @@ import functools
 import logging
 import os
 import typing
-from typing import Dict, List
+from typing import Any, Dict, List
 
 import torch
 import torch.distributed as dist
@@ -39,6 +39,23 @@ from nerfactory.viewer.server import viewer_utils
 logging.getLogger("PIL").setLevel(logging.WARNING)
 
 
+def train_loop(local_rank: int, world_size: int, config: Config) -> Any:
+    """Main training function that sets up and runs the trainer per process
+
+    Args:
+        local_rank (int): current rank of process
+        world_size (int): total number of gpus available
+        config (Config): config file specifying training regimen
+
+    Returns:
+        Any: TODO(): determine the return type
+    """
+    trainer = Trainer(config, local_rank, world_size)
+    trainer.setup()
+    trainer.train()
+    return 0
+
+
 class Trainer:
     """Trainer class
 
diff --git a/scripts/__init__.py b/scripts/__init__.py
deleted file mode 100644
index e69de29bb..000000000
diff --git a/scripts/run_train.py b/scripts/run_train.py
index 3e3dfb92c..edfde0ced 100644
--- a/scripts/run_train.py
+++ b/scripts/run_train.py
@@ -19,7 +19,7 @@ import yaml
 from hydra.core.config_store import ConfigStore
 from omegaconf import DictConfig
 
-from nerfactory.engine.trainer import Trainer
+from nerfactory.engine.trainer import train_loop
 from nerfactory.utils import comms, profiler
 from nerfactory.utils.config import Config, setup_config
 
@@ -112,23 +112,6 @@ def _distributed_worker(
     return output
 
 
-def _train(local_rank: int, world_size: int, config: Config) -> Any:
-    """Main training function that sets up and runs the trainer per process
-
-    Args:
-        local_rank (int): current rank of process
-        world_size (int): total number of gpus available
-        config (Config): config file specifying training regimen
-
-    Returns:
-        Any: TODO(): determine the return type
-    """
-    trainer = Trainer(config, local_rank, world_size)
-    trainer.setup()
-    trainer.train()
-    return 0
-
-
 def launch(
     main_func: Callable,
     num_gpus_per_machine: int,
@@ -219,7 +202,7 @@ def main(config: DictConfig):
         print(yaml.dump(unrolled_config, sort_keys=False, default_flow_style=False))
 
     launch(
-        _train,
+        train_loop,
         config.machine.num_gpus,
         num_machines=config.machine.num_machines,
         machine_rank=config.machine.machine_rank,
diff --git a/tests/test_run_train.py b/tests/test_run_train.py
index 83ac9d726..ffcbef041 100644
--- a/tests/test_run_train.py
+++ b/tests/test_run_train.py
@@ -7,9 +7,10 @@ import os
 import pytest
 from hydra import compose, initialize
 from hydra.core.global_hydra import GlobalHydra
-from omegaconf import DictConfig
+from omegaconf import DictConfig, open_dict
 
-from scripts.run_train import main
+from nerfactory.engine.trainer import train_loop
+from nerfactory.utils.config import setup_config
 
 BLACKLIST = ["graph_semantic_nerf.yaml", "graph_mipnerf_360.yaml", "graph_instant_ngp.yaml"]
 
@@ -29,11 +30,12 @@ def set_reduced_config(config: DictConfig):
     }
     config.pipeline.dataloader.train_num_images_to_sample_from = 1
     config.pipeline.dataloader.train_num_rays_per_batch = 4
-    config.pipeline.dataloader.eval_dataset = {
-        "_target_": "nerfactory.dataloaders.datasets.Blender",
-        "data_directory": "tests/data/lego_test",
-        "downscale_factor": 16,
-    }
+    with open_dict(config):
+        config.pipeline.dataloader.eval_dataset = {
+            "_target_": "nerfactory.dataloaders.datasets.Blender",
+            "data_directory": "tests/data/lego_test",
+            "downscale_factor": 16,
+        }
 
     # reduce model factors
     config.pipeline.model.num_coarse_samples = 4
@@ -61,7 +63,9 @@ def test_run_train():
         initialize(version_base="1.2", config_path="../configs/")  # relative to test path
         config = compose(config_path)
         config = set_reduced_config(config)
-        main(config)
+        config = setup_config(config)
+
+        train_loop(local_rank=0, world_size=0, config=config)
 
         GlobalHydra.instance().clear()