removing callable modules from script/

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
 

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,

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()