Integrating logic for setting up optimizers in run_train and config files (#33)

* added optimizers

* added optimizer logic for setting it up; todo-- fix optimizers and add step logic

* typing out the alphabet

.github/workflows/code_checks.yml

@@ -33,7 +33,7 @@ jobs:
           pylint mattport tests \
             --max-line-length=120 \
             --generated-members=numpy.*,torch.*,cv2.*,cv.* \
-            --good-names=e,c,i,j,x,y,lr,K,fx,fy,cx,cy,H,W,ts,t,u \
+            --good-names=e,c,i,j,x,y,lr,K,fx,fy,cx,cy,H,W,ts,t,u,k,v \
             --disable=too-many-arguments \
             --disable=too-many-locals \
             --disable=fixme

.vscode/settings.json

@@ -29,7 +29,7 @@
     "python.linting.pylintArgs": [
         "--generated-members=numpy.*,torch.*,cv2.*,cv.*",
         "--max-line-length=120",
-        "--good-names=e,c,i,j,x,y,lr,K,fx,fy,cx,cy,H,W,ts,t,u",
+        "--good-names=e,c,i,j,x,y,lr,K,fx,fy,cx,cy,H,W,ts,t,u,k,v",
         "--disable=too-many-arguments",
         "--disable=fixme",
         "--disable=too-many-locals"

configs/config.yaml

@@ -1,7 +1,7 @@
 # pretrained loading options
 resume_train:
-  load_dir: /tmp/mattport_models/
-  load_step: 10
+  load_dir: # pretrained model directory
+  load_step: # pretrained model's step number
 
 save_step: 10
 model_dir: /tmp/mattport_models/
@@ -20,8 +20,23 @@ dataloader:
   num_rays_per_batch: 1024
 
 # optimizer options
-optimizer:
-  lr: 1e-4
+param_groups:
+  graph:
+    optimizer:
+      type: Adam
+      lr: 1e-4
+    scheduler:
+      type: StepLR
+      step_size: 30
+      gamma: 0.1
+  camera:
+    optimizer:
+      type: SGD
+      lr: 1e-4
+    scheduler:
+      type: ExponentialLR 
+      gamma: 0.1
+
 
 # machine options
 machine_config:

mattport/nerf/graph/base.py

@@ -3,9 +3,10 @@ The Graph module contains all trainable parameters.
 """
 from abc import abstractmethod
 from dataclasses import dataclass
-from typing import Dict, Optional
+from typing import Dict, List, Optional
 
 from torch import nn
+from torch.nn import Parameter
 from torchtyping import TensorType
 
 
@@ -43,6 +44,10 @@ class Graph(nn.Module):
     def populate_modules(self):
         """Initializes the modules that are part of the network."""
 
+    @abstractmethod
+    def get_param_groups(self) -> Dict[str, List[Parameter]]:
+        """Returns the parameter groups of the network in a dictionary format"""
+
     def get_in_dim(self, curr_node: Node) -> None:
         """Dynamically calculates and sets the input dimensions of the modules based on dependency graph
 

mattport/nerf/graph/nerf.py

@@ -3,6 +3,8 @@ Implementation of vanilla nerf.
 """
 
 
+from typing import Dict, List
+from torch.nn import Parameter
 from torchtyping import TensorType
 
 from mattport.nerf.field.nerf import NeRFField
@@ -41,6 +43,17 @@ class NeRFGraph(Graph):
         # losses
         self.rgb_loss = MSELoss()
 
+    def get_param_groups(self) -> Dict[str, List[Parameter]]:
+        """Obtain the parameter groups for the optimizers
+
+        Returns:
+            Dict[str, List[Parameter]]: Mapping of different parameter groups
+        """
+        param_groups = {}
+        param_groups["camera"] = [dict(params=self.ray_generator.parameters())]
+        param_groups["graph"] = [dict(params=self.field_coarse.parameters()), dict(params=self.field_fine.parameters())]
+        return param_groups
+
     def forward(self, ray_indices: TensorType["num_rays", 3]):
         """Takes in the ray indices and renders out values."""
         # get the rays:

mattport/nerf/optimizers.py

@@ -18,18 +18,19 @@ class Optimizers:
             config (DictConfig): _description_
             param_dict (Dict[str, List[Parameter]]): _description_
         """
-        # TODO() add param_name `learning_rate`` and `weight_decay`, etc... to config file
         self.config = config
         self.optimizers = {}
         self.schedulers = {}
         for param_name, params in param_groups.items():
-            optimizer = getattr(torch.optim, config.param_name.optimizer.type)
-            kwargs = {k: v for k, v in config.param_name.optimizer.items() if k != "type"}
+            optimizer_config = config[param_name].optimizer
+            optimizer = getattr(torch.optim, optimizer_config.type)
+            kwargs = {k: v for k, v in optimizer_config.items() if k != "type"}
             self.optimizers[param_name] = optimizer(params, **kwargs)
-            if config.param_name.scheduler:
-                scheduler = getattr(torch.optim.lr_scheduler, config.param_name.scheduler.type)
-                kwargs = {k: v for k, v in config.param_name.scheduler.items() if k != "type"}
-                self.schedulers[param_name] = scheduler(optimizer, **kwargs)
+            if config[param_name].scheduler:
+                scheduler_config = config[param_name].scheduler
+                scheduler = getattr(torch.optim.lr_scheduler, scheduler_config.type)
+                kwargs = {k: v for k, v in scheduler_config.items() if k != "type"}
+                self.schedulers[param_name] = scheduler(self.optimizers[param_name], **kwargs)
 
     def optimizer_step(self, param_name: str) -> None:
         """fetch and step corresponding optimizer

mattport/nerf/trainer.py

@@ -8,12 +8,12 @@ import torch
 import torch.distributed as dist
 from hydra.utils import instantiate
 from omegaconf import DictConfig
-from torch import optim
 from torch.nn.parallel import DistributedDataParallel as DDP
 from torch.utils.data import DataLoader
 
 from mattport.nerf.dataset.image_dataset import ImageDataset, collate_batch
 from mattport.nerf.dataset.utils import get_dataset_inputs
+from mattport.nerf.optimizers import Optimizers
 from mattport.utils.decorators import check_main_thread
 from mattport.utils.writer import LocalWriter
 
@@ -36,7 +36,7 @@ class Trainer:
         self.test_dataloader = None
         # model variables
         self.graph = None
-        self.optimizer = None
+        self.optimizers = None
         self.start_step = 0
         # logging variables
         self.is_main_thread = local_rank % world_size == 0
@@ -65,18 +65,18 @@ class Trainer:
         self.graph = instantiate(
             self.config.graph, intrinsics=dataset_inputs.intrinsics, camera_to_world=dataset_inputs.camera_to_world
         ).to(f"cuda:{self.local_rank}")
-        self.setup_optimizer()
+        self.setup_optimizers()
 
-        if self.config.resume_train:
+        if self.config.resume_train.load_dir:
             self.load_checkpoint(self.config.resume_train)
 
         if self.world_size > 1:
             self.graph = DDP(self.graph, device_ids=[self.local_rank])
             dist.barrier(device_ids=[self.local_rank])
 
-    def setup_optimizer(self):
+    def setup_optimizers(self):
         """_summary_"""
-        self.optimizer = optim.Adam(self.graph.parameters(), lr=self.config.optimizer.lr)
+        self.optimizers = Optimizers(self.config.param_groups, self.graph.get_param_groups())
 
     def load_checkpoint(self, load_config: DictConfig) -> int:
         """Load the checkpoint from the given path
@@ -91,19 +91,20 @@ class Trainer:
         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()})
-        self.optimizer.load_state_dict(loaded_state["optimizer"])
+        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 optimizer
+        """Save the model and optimizers
 
         Args:
             output_dir (str): directory to save the checkpoint
             step (int): number of steps in training for given checkpoint
             model (Graph): Graph model to be saved
-            optimizer (Optimizer): Optimizer to be saved
+            optimizers (Optimizers): Optimizers to be saved
         """
         if not os.path.exists(output_dir):
             os.makedirs(output_dir)
@@ -112,7 +113,7 @@ class Trainer:
             {
                 "step": step,
                 "model": self.graph.module.state_dict() if hasattr(self.graph, "module") else self.graph.state_dict(),
-                "optimizer": self.optimizer.state_dict(),
+                "optimizers": {k: v.state_dict() for (k, v) in self.optimizers.optimizers.items()},
             },
             ckpt_path,
         )