From 17d837babfcd923e6584338a1f35960ced0cd74a Mon Sep 17 00:00:00 2001 From: evonneng Date: Thu, 21 Apr 2022 13:23:18 -0700 Subject: [PATCH] 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 | 2 +- .vscode/settings.json | 2 +- configs/config.yaml | 23 +++++++++++++++++++---- mattport/nerf/graph/base.py | 7 ++++++- mattport/nerf/graph/nerf.py | 13 +++++++++++++ mattport/nerf/optimizers.py | 15 ++++++++------- mattport/nerf/trainer.py | 21 +++++++++++---------- 7 files changed, 59 insertions(+), 24 deletions(-) diff --git a/.github/workflows/code_checks.yml b/.github/workflows/code_checks.yml index af536cb09..270a67066 100644 --- a/.github/workflows/code_checks.yml +++ b/.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 diff --git a/.vscode/settings.json b/.vscode/settings.json index ab9efac9c..80164259e 100644 --- a/.vscode/settings.json +++ b/.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" diff --git a/configs/config.yaml b/configs/config.yaml index 9e2847eea..41323f313 100644 --- a/configs/config.yaml +++ b/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: diff --git a/mattport/nerf/graph/base.py b/mattport/nerf/graph/base.py index 7da490633..2bce70e44 100644 --- a/mattport/nerf/graph/base.py +++ b/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 diff --git a/mattport/nerf/graph/nerf.py b/mattport/nerf/graph/nerf.py index 4c2598704..4f7a3b397 100644 --- a/mattport/nerf/graph/nerf.py +++ b/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: diff --git a/mattport/nerf/optimizers.py b/mattport/nerf/optimizers.py index 641e0e632..e88b8e442 100644 --- a/mattport/nerf/optimizers.py +++ b/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 diff --git a/mattport/nerf/trainer.py b/mattport/nerf/trainer.py index a8e866e98..3859e2a38 100644 --- a/mattport/nerf/trainer.py +++ b/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, )