diff --git a/.vscode/settings.json b/.vscode/settings.json index 8b0f1d750..1bd31d744 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -34,7 +34,7 @@ "python.envFile": "${workspaceFolder}/.env", "python.formatting.provider": "none", "black-formatter.args": ["--line-length=120"], - "python.linting.pylintEnabled": false, + "python.linting.pylintEnabled": true, "python.linting.flake8Enabled": false, "python.linting.enabled": true, "python.testing.unittestEnabled": false, diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 3e5939eed..55503a720 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -71,7 +71,7 @@ from nerfstudio.models.semantic_nerfw import SemanticNerfWModelConfig from nerfstudio.models.tensorf import TensoRFModelConfig from nerfstudio.models.vanilla_nerf import NeRFModel, VanillaModelConfig from nerfstudio.pipelines.base_pipeline import VanillaPipelineConfig -from nerfstudio.data.datamanagers.full_images_datamanager import FullImageDatamanagerConfig +from nerfstudio.data.datamanagers.full_images_datamanager import FullImageDatamanagerConfig,FullImageDatamanager from nerfstudio.pipelines.dynamic_batch import DynamicBatchPipelineConfig from nerfstudio.plugins.registry import discover_methods @@ -597,17 +597,16 @@ method_configs["neus-facto"] = TrainerConfig( viewer=ViewerConfig(num_rays_per_chunk=1 << 15), vis="viewer", ) -from dataclasses import field -from collections import defaultdict + method_configs["gaussian-splatting"] = TrainerConfig( method_name="gaussian-splatting", - steps_per_eval_image=10, - steps_per_eval_batch=10, + steps_per_eval_image=100, + steps_per_eval_batch=100, steps_per_save=2000, - steps_per_eval_all_images=1000000, # set to a very large model so we don't eval with all images + steps_per_eval_all_images=100000, max_num_iterations=30000, mixed_precision=False, - gradient_accumulation_steps = {'camera_opt': 100}, + gradient_accumulation_steps = {'camera_opt': 100,'color':10,'shs':10}, pipeline=VanillaPipelineConfig( datamanager=FullImageDatamanagerConfig( dataparser=ColmapDataParserConfig(load_3D_points=True), @@ -642,7 +641,7 @@ method_configs["gaussian-splatting"] = TrainerConfig( }, "scaling": { "optimizer": AdamOptimizerConfig(lr=0.005, eps=1e-15), - "scheduler": None + "scheduler": ExponentialDecaySchedulerConfig(lr_final=1e-3, max_steps=30000) }, "rotation": { "optimizer": AdamOptimizerConfig(lr=0.001, eps=1e-15), diff --git a/nerfstudio/data/datamanagers/full_images_datamanager.py b/nerfstudio/data/datamanagers/full_images_datamanager.py index 31dff3da1..16de6f3f0 100644 --- a/nerfstudio/data/datamanagers/full_images_datamanager.py +++ b/nerfstudio/data/datamanagers/full_images_datamanager.py @@ -40,6 +40,7 @@ from nerfstudio.data.datamanagers.base_datamanager import DataManager, DataManag from nerfstudio.data.dataparsers.base_dataparser import DataparserOutputs from nerfstudio.data.dataparsers.nerfstudio_dataparser import NerfstudioDataParserConfig from nerfstudio.data.datasets.base_dataset import InputDataset +# from nerfstudio.data.utils.dataloaders import FixedIndicesEvalDataloader from nerfstudio.utils.misc import get_orig_class from nerfstudio.utils.rich_utils import CONSOLE @@ -146,6 +147,10 @@ class FullImageDatamanager(DataManager, Generic[TDataset]): # crop the image and update the intrinsics accordingly x, y, w, h = roi image = image[y : y + h, x : x + w] + if 'mask' in data: + data['mask'] = data['mask'][y : y + h, x : x + w] + if 'depth_image' in data: + data['depth_image'] = data['depth_image'][y : y + h, x : x + w] K = newK # update the width, height self.train_dataset.cameras.width[i] = w @@ -215,8 +220,8 @@ class FullImageDatamanager(DataManager, Generic[TDataset]): image = image[y : y + h, x : x + w] K = newK # update the width, height - self.train_dataset.cameras.width[i] = w - self.train_dataset.cameras.height[i] = h + self.eval_dataset.cameras.width[i] = w + self.eval_dataset.cameras.height[i] = h elif camera.camera_type.item() == CameraType.FISHEYE.value: distortion_params = np.array( [distortion_params[0], distortion_params[1], distortion_params[2], distortion_params[3]] @@ -250,7 +255,6 @@ class FullImageDatamanager(DataManager, Generic[TDataset]): self.eval_dataset.cameras.fy[i] = float(K[1, 1]) self.eval_dataset.cameras.cx[i] = float(K[0, 2]) self.eval_dataset.cameras.cy[i] = float(K[1, 2]) - # print(self.eval_dataset.cameras[i].reshape((1,))) if cache_images_option == "gpu": for cache in cached_train: @@ -311,7 +315,7 @@ class FullImageDatamanager(DataManager, Generic[TDataset]): def setup_eval(self): """Sets up the data loader for evaluation""" - + def get_param_groups(self) -> Dict[str, List[Parameter]]: """Get the param groups for the data manager. Returns: @@ -350,13 +354,10 @@ class FullImageDatamanager(DataManager, Generic[TDataset]): # Make sure to re-populate the unseen cameras list if we have exhausted it if len(self.eval_unseen_cameras) == 0: self.eval_unseen_cameras = [i for i in range(len(self.eval_dataset))] - data = self.eval_dataset.get_data(image_idx) + data = deepcopy(self.cached_eval[image_idx]) data["image"] = data["image"].to(self.device) assert len(self.eval_dataset.cameras.shape) == 1, "Assumes single batch dimension" camera = self.eval_dataset.cameras[image_idx : image_idx + 1].to(self.device) - # TODO: camera shape torch.Size([]) for eval dataset - # should be torch.Size([1]) - print(camera.shape, data.keys()) return camera, data def next_eval_image(self, step: int) -> Tuple[int, Cameras, Dict]: @@ -369,7 +370,7 @@ class FullImageDatamanager(DataManager, Generic[TDataset]): # Make sure to re-populate the unseen cameras list if we have exhausted it if len(self.eval_unseen_cameras) == 0: self.eval_unseen_cameras = [i for i in range(len(self.eval_dataset))] - data = self.eval_dataset.get_data(image_idx) + data = deepcopy(self.cached_eval[image_idx]) data["image"] = data["image"].to(self.device) assert len(self.eval_dataset.cameras.shape) == 1, "Assumes single batch dimension" camera = self.eval_dataset.cameras[image_idx: image_idx + 1].to(self.device) diff --git a/nerfstudio/engine/hacked_adam.py b/nerfstudio/engine/hacked_adam.py new file mode 100644 index 000000000..6ddfa5e4a --- /dev/null +++ b/nerfstudio/engine/hacked_adam.py @@ -0,0 +1,583 @@ +from typing import List, Optional + +import torch +from torch import Tensor +from torch.optim.optimizer import (Optimizer, _use_grad_for_differentiable, _get_value, _stack_if_compiling, + _dispatch_sqrt, _default_to_fused_or_foreach, _capturable_doc, + _differentiable_doc, _foreach_doc, _fused_doc, _maximize_doc) +from torch.utils._foreach_utils import _group_tensors_by_device_and_dtype + +__all__ = ['Adam', 'adam'] + + +class Adam(Optimizer): + def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, + weight_decay=0, amsgrad=False, *, foreach: Optional[bool] = None, + maximize: bool = False, capturable: bool = False, + differentiable: bool = False, fused: Optional[bool] = None): + if not 0.0 <= lr: + raise ValueError("Invalid learning rate: {}".format(lr)) + if not 0.0 <= eps: + raise ValueError("Invalid epsilon value: {}".format(eps)) + if not 0.0 <= betas[0] < 1.0: + raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0])) + if not 0.0 <= betas[1] < 1.0: + raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1])) + if not 0.0 <= weight_decay: + raise ValueError("Invalid weight_decay value: {}".format(weight_decay)) + + defaults = dict(lr=lr, betas=betas, eps=eps, + weight_decay=weight_decay, amsgrad=amsgrad, + maximize=maximize, foreach=foreach, capturable=capturable, + differentiable=differentiable, fused=fused) + super().__init__(params, defaults) + self.gauss_scaling = 1.0 + + if fused: + if differentiable: + raise RuntimeError("`fused` does not support `differentiable`") + self._step_supports_amp_scaling = True + # TODO(crcrpar): [low prec params & their higher prec copy] + # Suppor AMP with FP16/BF16 model params which would need + # higher prec copy of params to do update math in higher prec to + # alleviate the loss of information. + if not all( + p.is_cuda and torch.is_floating_point(p) + for pg in self.param_groups for p in pg['params'] + ): + raise RuntimeError("`fused=True` requires all the params to be CUDA, floating point Tensor") + if foreach: + raise RuntimeError("`fused` and `foreach` cannot be `True` together.") + + def __setstate__(self, state): + super().__setstate__(state) + for group in self.param_groups: + group.setdefault('amsgrad', False) + group.setdefault('maximize', False) + group.setdefault('foreach', None) + group.setdefault('capturable', False) + group.setdefault('differentiable', False) + group.setdefault('fused', None) + state_values = list(self.state.values()) + step_is_tensor = (len(state_values) != 0) and torch.is_tensor(state_values[0]['step']) + if not step_is_tensor: + for s in state_values: + s['step'] = torch.tensor(float(s['step'])) + + def _init_group( + self, + group, + params_with_grad, + grads, + exp_avgs, + exp_avg_sqs, + max_exp_avg_sqs, + state_steps + ): + for p in group['params']: + if p.grad is not None: + params_with_grad.append(p) + if p.grad.is_sparse: + raise RuntimeError('Adam does not support sparse gradients, please consider SparseAdam instead') + grads.append(p.grad) + + state = self.state[p] + # Lazy state initialization + if len(state) == 0: + state['step'] = ( + torch.zeros((1,), dtype=torch.float, device=p.device) + if group['capturable'] or group['fused'] + else torch.tensor(0.) + ) + # Exponential moving average of gradient values + state['exp_avg'] = torch.zeros_like(p, memory_format=torch.preserve_format) + # Exponential moving average of squared gradient values + state['exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format) + if group['amsgrad']: + # Maintains max of all exp. moving avg. of sq. grad. values + state['max_exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format) + + exp_avgs.append(state['exp_avg']) + exp_avg_sqs.append(state['exp_avg_sq']) + + if group['amsgrad']: + max_exp_avg_sqs.append(state['max_exp_avg_sq']) + if group['differentiable'] and state['step'].requires_grad: + raise RuntimeError('`requires_grad` is not supported for `step` in differentiable mode') + state_steps.append(state['step']) + + @_use_grad_for_differentiable + def step(self, closure=None): + """Performs a single optimization step. + + Args: + closure (Callable, optional): A closure that reevaluates the model + and returns the loss. + """ + self._cuda_graph_capture_health_check() + + loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() + + for group in self.param_groups: + params_with_grad = [] + grads = [] + exp_avgs = [] + exp_avg_sqs = [] + max_exp_avg_sqs = [] + state_steps = [] + beta1, beta2 = group['betas'] + + self._init_group( + group, + params_with_grad, + grads, + exp_avgs, + exp_avg_sqs, + max_exp_avg_sqs, + state_steps) + + adam( + params_with_grad, + grads, + exp_avgs, + exp_avg_sqs, + max_exp_avg_sqs, + state_steps, + amsgrad=group['amsgrad'], + beta1=beta1, + beta2=beta2, + lr=group['lr'], + weight_decay=group['weight_decay'], + eps=group['eps'], + maximize=group['maximize'], + foreach=group['foreach'], + capturable=group['capturable'], + differentiable=group['differentiable'], + fused=group['fused'], + grad_scale=getattr(self, "grad_scale", None), + found_inf=getattr(self, "found_inf", None), + gauss_scaling=self.gauss_scaling + ) + + return loss + + +Adam.__doc__ = r"""Implements Adam algorithm. + + .. math:: + \begin{aligned} + &\rule{110mm}{0.4pt} \\ + &\textbf{input} : \gamma \text{ (lr)}, \beta_1, \beta_2 + \text{ (betas)},\theta_0 \text{ (params)},f(\theta) \text{ (objective)} \\ + &\hspace{13mm} \lambda \text{ (weight decay)}, \: \textit{amsgrad}, + \:\textit{maximize} \\ + &\textbf{initialize} : m_0 \leftarrow 0 \text{ ( first moment)}, + v_0\leftarrow 0 \text{ (second moment)},\: \widehat{v_0}^{max}\leftarrow 0\\[-1.ex] + &\rule{110mm}{0.4pt} \\ + &\textbf{for} \: t=1 \: \textbf{to} \: \ldots \: \textbf{do} \\ + + &\hspace{5mm}\textbf{if} \: \textit{maximize}: \\ + &\hspace{10mm}g_t \leftarrow -\nabla_{\theta} f_t (\theta_{t-1}) \\ + &\hspace{5mm}\textbf{else} \\ + &\hspace{10mm}g_t \leftarrow \nabla_{\theta} f_t (\theta_{t-1}) \\ + &\hspace{5mm}\textbf{if} \: \lambda \neq 0 \\ + &\hspace{10mm} g_t \leftarrow g_t + \lambda \theta_{t-1} \\ + &\hspace{5mm}m_t \leftarrow \beta_1 m_{t-1} + (1 - \beta_1) g_t \\ + &\hspace{5mm}v_t \leftarrow \beta_2 v_{t-1} + (1-\beta_2) g^2_t \\ + &\hspace{5mm}\widehat{m_t} \leftarrow m_t/\big(1-\beta_1^t \big) \\ + &\hspace{5mm}\widehat{v_t} \leftarrow v_t/\big(1-\beta_2^t \big) \\ + &\hspace{5mm}\textbf{if} \: amsgrad \\ + &\hspace{10mm}\widehat{v_t}^{max} \leftarrow \mathrm{max}(\widehat{v_t}^{max}, + \widehat{v_t}) \\ + &\hspace{10mm}\theta_t \leftarrow \theta_{t-1} - \gamma \widehat{m_t}/ + \big(\sqrt{\widehat{v_t}^{max}} + \epsilon \big) \\ + &\hspace{5mm}\textbf{else} \\ + &\hspace{10mm}\theta_t \leftarrow \theta_{t-1} - \gamma \widehat{m_t}/ + \big(\sqrt{\widehat{v_t}} + \epsilon \big) \\ + &\rule{110mm}{0.4pt} \\[-1.ex] + &\bf{return} \: \theta_t \\[-1.ex] + &\rule{110mm}{0.4pt} \\[-1.ex] + \end{aligned} + + For further details regarding the algorithm we refer to `Adam: A Method for Stochastic Optimization`_. + """ + r""" + Args: + params (iterable): iterable of parameters to optimize or dicts defining + parameter groups + lr (float, optional): learning rate (default: 1e-3) + betas (Tuple[float, float], optional): coefficients used for computing + running averages of gradient and its square (default: (0.9, 0.999)) + eps (float, optional): term added to the denominator to improve + numerical stability (default: 1e-8) + weight_decay (float, optional): weight decay (L2 penalty) (default: 0) + amsgrad (bool, optional): whether to use the AMSGrad variant of this + algorithm from the paper `On the Convergence of Adam and Beyond`_ + (default: False) + {foreach} + {maximize} + {capturable} + {differentiable} + {fused} + .. _Adam\: A Method for Stochastic Optimization: + https://arxiv.org/abs/1412.6980 + .. _On the Convergence of Adam and Beyond: + https://openreview.net/forum?id=ryQu7f-RZ + + """.format(foreach=_foreach_doc, maximize=_maximize_doc, capturable=_capturable_doc, + differentiable=_differentiable_doc, fused=_fused_doc) + + +def adam(params: List[Tensor], + grads: List[Tensor], + exp_avgs: List[Tensor], + exp_avg_sqs: List[Tensor], + max_exp_avg_sqs: List[Tensor], + state_steps: List[Tensor], + # kwonly args with defaults are not supported by functions compiled with torchscript issue #70627 + # setting this as kwarg for now as functional API is compiled by torch/distributed/optim + foreach: Optional[bool] = None, + capturable: bool = False, + differentiable: bool = False, + fused: Optional[bool] = None, + grad_scale: Optional[Tensor] = None, + found_inf: Optional[Tensor] = None, + *, + amsgrad: bool, + beta1: float, + beta2: float, + lr: float, + weight_decay: float, + eps: float, + maximize: bool, + gauss_scaling: float = 1.0 + ): + r"""Functional API that performs Adam algorithm computation. + See :class:`~torch.optim.Adam` for details. + """ + + # Respect when the user inputs False/True for foreach or fused. We only want to change + # the default when neither have been user-specified. Note that we default to foreach + # and pass False to use_fused. This is not a mistake--we want to give the fused impl + # bake-in time before making it the default, even if it is typically faster. + if fused is None and foreach is None: + _, foreach = _default_to_fused_or_foreach(params, differentiable, use_fused=False) + if fused is None: + fused = False + if foreach is None: + foreach = False + + if not all(isinstance(t, torch.Tensor) for t in state_steps): + raise RuntimeError("API has changed, `state_steps` argument must contain a list of singleton tensors") + + if foreach and torch.jit.is_scripting(): + raise RuntimeError('torch.jit.script not supported with foreach optimizers') + + if fused and not torch.jit.is_scripting(): + func = _fused_adam + elif foreach and not torch.jit.is_scripting(): + func = _multi_tensor_adam + else: + func = _single_tensor_adam + + func(params, + grads, + exp_avgs, + exp_avg_sqs, + max_exp_avg_sqs, + state_steps, + amsgrad=amsgrad, + beta1=beta1, + beta2=beta2, + lr=lr, + weight_decay=weight_decay, + eps=eps, + maximize=maximize, + capturable=capturable, + differentiable=differentiable, + grad_scale=grad_scale, + found_inf=found_inf, + gauss_scaling=gauss_scaling + ) + + +def _single_tensor_adam(params: List[Tensor], + grads: List[Tensor], + exp_avgs: List[Tensor], + exp_avg_sqs: List[Tensor], + max_exp_avg_sqs: List[Tensor], + state_steps: List[Tensor], + grad_scale: Optional[Tensor], + found_inf: Optional[Tensor], + *, + amsgrad: bool, + beta1: float, + beta2: float, + lr: float, + weight_decay: float, + eps: float, + maximize: bool, + capturable: bool, + differentiable: bool): + + assert grad_scale is None and found_inf is None + + for i, param in enumerate(params): + + grad = grads[i] if not maximize else -grads[i] + exp_avg = exp_avgs[i] + exp_avg_sq = exp_avg_sqs[i] + step_t = state_steps[i] + + if capturable: + assert param.is_cuda and step_t.is_cuda, "If capturable=True, params and state_steps must be CUDA tensors." + + # update step + step_t += 1 + + if weight_decay != 0: + grad = grad.add(param, alpha=weight_decay) + + if torch.is_complex(param): + grad = torch.view_as_real(grad) + exp_avg = torch.view_as_real(exp_avg) + exp_avg_sq = torch.view_as_real(exp_avg_sq) + param = torch.view_as_real(param) + + # Decay the first and second moment running average coefficient + exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1) + exp_avg_sq.mul_(beta2).addcmul_(grad, grad.conj(), value=1 - beta2) + + if capturable or differentiable: + step = step_t + + # 1 - beta1 ** step can't be captured in a CUDA graph, even if step is a CUDA tensor + # (incurs "RuntimeError: CUDA error: operation not permitted when stream is capturing") + bias_correction1 = 1 - torch.pow(beta1, step) + bias_correction2 = 1 - torch.pow(beta2, step) + + step_size = lr / bias_correction1 + step_size_neg = step_size.neg() + + bias_correction2_sqrt = bias_correction2.sqrt() + + if amsgrad: + # Maintains the maximum of all 2nd moment running avg. till now + if differentiable: + max_exp_avg_sqs_i = max_exp_avg_sqs[i].clone() + else: + max_exp_avg_sqs_i = max_exp_avg_sqs[i] + max_exp_avg_sqs[i].copy_(torch.maximum(max_exp_avg_sqs_i, exp_avg_sq)) + # Uses the max. for normalizing running avg. of gradient + # Folds in (admittedly ugly) 1-elem step_size math here to avoid extra param-set-sized read+write + # (can't fold it into addcdiv_ below because addcdiv_ requires value is a Number, not a Tensor) + denom = (max_exp_avg_sqs[i].sqrt() / (bias_correction2_sqrt * step_size_neg)).add_(eps / step_size_neg) + else: + denom = (exp_avg_sq.sqrt() / (bias_correction2_sqrt * step_size_neg)).add_(eps / step_size_neg) + param.addcdiv_(exp_avg, denom) + else: + step = _get_value(step_t) + + bias_correction1 = 1 - beta1 ** step + bias_correction2 = 1 - beta2 ** step + + step_size = lr / bias_correction1 + + bias_correction2_sqrt = _dispatch_sqrt(bias_correction2) + + if amsgrad: + # Maintains the maximum of all 2nd moment running avg. till now + torch.maximum(max_exp_avg_sqs[i], exp_avg_sq, out=max_exp_avg_sqs[i]) + # Use the max. for normalizing running avg. of gradient + denom = (max_exp_avg_sqs[i].sqrt() / bias_correction2_sqrt).add_(eps) + else: + denom = (exp_avg_sq.sqrt() / bias_correction2_sqrt).add_(eps) + param.addcdiv_(exp_avg, denom, value=-step_size) + + +def _multi_tensor_adam(params: List[Tensor], + grads: List[Tensor], + exp_avgs: List[Tensor], + exp_avg_sqs: List[Tensor], + max_exp_avg_sqs: List[Tensor], + state_steps: List[Tensor], + grad_scale: Optional[Tensor], + found_inf: Optional[Tensor], + *, + amsgrad: bool, + beta1: float, + beta2: float, + lr: float, + weight_decay: float, + eps: float, + maximize: bool, + capturable: bool, + differentiable: bool, + gauss_scaling: float = 1.0): + if len(params) == 0: + return + + if capturable: + assert all(p.is_cuda and step.is_cuda for p, step in zip(params, state_steps)), \ + "If capturable=True, params and state_steps must be CUDA tensors." + + assert grad_scale is None and found_inf is None + + assert not differentiable, "_foreach ops don't support autograd" + + grouped_tensors = _group_tensors_by_device_and_dtype([params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, state_steps]) + for (device_params, device_grads, device_exp_avgs, device_exp_avg_sqs, + device_max_exp_avg_sqs, device_state_steps) in grouped_tensors.values(): + + if maximize: + device_grads = torch._foreach_neg(tuple(device_grads)) # type: ignore[assignment] + + # Handle complex parameters + device_grads = [torch.view_as_real(x) if torch.is_complex(x) else x for x in device_grads] + device_exp_avgs = [torch.view_as_real(x) if torch.is_complex(x) else x for x in device_exp_avgs] + device_exp_avg_sqs = [torch.view_as_real(x) if torch.is_complex(x) else x for x in device_exp_avg_sqs] + params_ = [torch.view_as_real(x) if torch.is_complex(x) else x for x in device_params] + + # update steps + torch._foreach_add_(device_state_steps, 1) + + if weight_decay != 0: + device_grads = torch._foreach_add(device_grads, device_params, alpha=weight_decay) + + # Decay the first and second moment running average coefficient + torch._foreach_mul_(device_exp_avgs, beta1) + torch._foreach_add_(device_exp_avgs, device_grads, alpha=1 - beta1) + + torch._foreach_mul_(device_exp_avg_sqs, beta2) + torch._foreach_addcmul_(device_exp_avg_sqs, device_grads, device_grads, 1 - beta2) + + if capturable: + # TODO: use foreach_pow if/when foreach_pow is added + bias_correction1 = [torch.pow(beta1, step) for step in device_state_steps] + bias_correction2 = [torch.pow(beta2, step) for step in device_state_steps] + # foreach_sub doesn't allow a scalar as the first arg + torch._foreach_sub_(bias_correction1, 1) + torch._foreach_sub_(bias_correction2, 1) + torch._foreach_neg_(bias_correction1) + torch._foreach_neg_(bias_correction2) + + # foreach_div doesn't allow a scalar as the first arg + step_size = torch._foreach_div(bias_correction1, lr) + torch._foreach_reciprocal_(step_size) + torch._foreach_neg_(step_size) + + bias_correction2_sqrt = torch._foreach_sqrt(bias_correction2) + + if amsgrad: + # Maintains the maximum of all 2nd moment running avg. till now + torch._foreach_maximum_(device_max_exp_avg_sqs, device_exp_avg_sqs) # type: ignore[assignment] + + # Use the max. for normalizing running avg. of gradient + max_exp_avg_sq_sqrt = torch._foreach_sqrt(device_max_exp_avg_sqs) + # Folds in (admittedly ugly) 1-elem step_size math here to avoid extra param-set-sized read+write + # (can't fold it into addcdiv_ below because addcdiv_ requires value is a Number, not a Tensor) + torch._foreach_div_(max_exp_avg_sq_sqrt, torch._foreach_mul(bias_correction2_sqrt, step_size)) + eps_over_step_size = torch._foreach_div(step_size, eps) + torch._foreach_reciprocal_(eps_over_step_size) + denom = torch._foreach_add(max_exp_avg_sq_sqrt, eps_over_step_size) + else: + exp_avg_sq_sqrt = torch._foreach_sqrt(device_exp_avg_sqs) + torch._foreach_div_(exp_avg_sq_sqrt, torch._foreach_mul(bias_correction2_sqrt, step_size)) + eps_over_step_size = torch._foreach_div(step_size, eps) + torch._foreach_reciprocal_(eps_over_step_size) + denom = torch._foreach_add(exp_avg_sq_sqrt, eps_over_step_size) + assert False + torch._foreach_addcdiv_(params_, device_exp_avgs, denom) + else: + bias_correction1 = [1 - beta1 ** _get_value(step) for step in device_state_steps] + bias_correction2 = [1 - beta2 ** _get_value(step) for step in device_state_steps] + + step_size = _stack_if_compiling([(lr / bc) * -1 for bc in bias_correction1]) + + bias_correction2_sqrt = [_dispatch_sqrt(bc) for bc in bias_correction2] + + if amsgrad: + # Maintains the maximum of all 2nd moment running avg. till now + torch._foreach_maximum_(device_max_exp_avg_sqs, device_exp_avg_sqs) + + # Use the max. for normalizing running avg. of gradient + max_exp_avg_sq_sqrt = torch._foreach_sqrt(device_max_exp_avg_sqs) + torch._foreach_div_(max_exp_avg_sq_sqrt, bias_correction2_sqrt) + denom = torch._foreach_add(max_exp_avg_sq_sqrt, eps) + else: + exp_avg_sq_sqrt = torch._foreach_sqrt(device_exp_avg_sqs) + torch._foreach_div_(exp_avg_sq_sqrt, bias_correction2_sqrt) + denom = torch._foreach_add(exp_avg_sq_sqrt, eps) + #apply the scaling here + exp_abs = device_exp_avgs[0].abs() + scaling = torch.ones_like(exp_abs) + biggers = (exp_abs>torch.tensor(gauss_scaling)).squeeze() + scaling[biggers] = (torch.tensor(gauss_scaling)/exp_abs)[biggers] + # scaling = torch.tensor(gauss_scaling) + device_exp_avgs = [x*scaling for i,x in enumerate(device_exp_avgs)] + torch._foreach_addcdiv_(params_, device_exp_avgs, denom, step_size) + + +def _fused_adam( + params: List[Tensor], + grads: List[Tensor], + exp_avgs: List[Tensor], + exp_avg_sqs: List[Tensor], + max_exp_avg_sqs: List[Tensor], + state_steps: List[Tensor], + grad_scale: Optional[Tensor], + found_inf: Optional[Tensor], + *, + amsgrad: bool, + beta1: float, + beta2: float, + lr: float, + weight_decay: float, + eps: float, + maximize: bool, + capturable: bool, # Needed for consistency. + differentiable: bool, +) -> None: + grad_scale_dict = {grad_scale.device: grad_scale} if grad_scale is not None else None + found_inf_dict = {found_inf.device: found_inf} if found_inf is not None else None + grouped_tensors = _group_tensors_by_device_and_dtype([params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, state_steps]) + for (device, dtype) in grouped_tensors: + ( + device_params, + device_grads, + device_exp_avgs, + device_exp_avg_sqs, + device_max_exp_avg_sqs, + device_state_steps, + ) = grouped_tensors[(device, dtype)] + device_grad_scale, device_found_inf = None, None + if grad_scale is not None: + if device not in grad_scale_dict: + grad_scale_dict[device] = grad_scale.to(device, non_blocking=True) + device_grad_scale = grad_scale_dict[device] + if found_inf is not None: + if found_inf not in found_inf_dict: + found_inf_dict[device] = found_inf.to(device, non_blocking=True) + device_found_inf = found_inf_dict[device] + torch._foreach_add_(device_state_steps, 1) + torch._fused_adam_( + device_params, + device_grads, + device_exp_avgs, + device_exp_avg_sqs, + device_max_exp_avg_sqs, + device_state_steps, + amsgrad=amsgrad, + lr=lr, + beta1=beta1, + beta2=beta2, + weight_decay=weight_decay, + eps=eps, + maximize=maximize, + grad_scale=device_grad_scale, + found_inf=device_found_inf, + ) + if device_found_inf is not None: + torch._foreach_sub_(device_state_steps, [device_found_inf] * len(device_state_steps)) diff --git a/nerfstudio/engine/optimizers.py b/nerfstudio/engine/optimizers.py index 5e35808c8..a1dc1c58f 100644 --- a/nerfstudio/engine/optimizers.py +++ b/nerfstudio/engine/optimizers.py @@ -27,7 +27,6 @@ from torch.nn.parameter import Parameter from nerfstudio.configs import base_config from nerfstudio.utils import writer - # Optimizer related configs @dataclass class OptimizerConfig(base_config.PrintableConfig): diff --git a/nerfstudio/engine/trainer.py b/nerfstudio/engine/trainer.py index a8da3f935..6cfde69bc 100644 --- a/nerfstudio/engine/trainer.py +++ b/nerfstudio/engine/trainer.py @@ -466,16 +466,6 @@ class Trainer: needs_zero = [group for group in self.optimizers.parameters.keys() if step % self.gradient_accumulation_steps[group] == 0] self.optimizers.zero_grad_some(needs_zero) cpu_or_cuda_str: str = self.device.split(":")[0] - # assert ( - # self.gradient_accumulation_steps > 0 - # ), f"gradient_accumulation_steps must be > 0, not {self.gradient_accumulation_steps}" - # for _ in range(self.gradient_accumulation_steps): - # with torch.autocast(device_type=cpu_or_cuda_str, enabled=self.mixed_precision): - # _, loss_dict, metrics_dict = self.pipeline.get_train_loss_dict(step=step) - # loss = functools.reduce(torch.add, loss_dict.values()) - # loss /= self.gradient_accumulation_steps - # self.grad_scaler.scale(loss).backward() # type: ignore - # self.optimizers.optimizer_scaler_step_all(self.grad_scaler) with torch.autocast(device_type=cpu_or_cuda_str, enabled=self.mixed_precision): _, loss_dict, metrics_dict = self.pipeline.get_train_loss_dict(step=step) diff --git a/nerfstudio/model_components/losses.py b/nerfstudio/model_components/losses.py index ff7ab1774..296e91520 100644 --- a/nerfstudio/model_components/losses.py +++ b/nerfstudio/model_components/losses.py @@ -591,6 +591,10 @@ def depth_ranking_loss(rendered_depth, gt_depth): and rendered_depth are from pixels with a radius of each other """ m = 1e-4 + if rendered_depth.shape[0] % 2 != 0: + # chop off one index + rendered_depth = rendered_depth[:-1, :] + gt_depth = gt_depth[:-1, :] dpt_diff = gt_depth[::2, :] - gt_depth[1::2, :] out_diff = rendered_depth[::2, :] - rendered_depth[1::2, :] + m differing_signs = torch.sign(dpt_diff) != torch.sign(out_diff) diff --git a/nerfstudio/models/gaussian_splatting.py b/nerfstudio/models/gaussian_splatting.py index fc6070b54..ad32c7f44 100644 --- a/nerfstudio/models/gaussian_splatting.py +++ b/nerfstudio/models/gaussian_splatting.py @@ -25,7 +25,7 @@ from nerfstudio.data.scene_box import OrientedBox from copy import deepcopy from nerfstudio.cameras.rays import RayBundle from torch.nn import Parameter -from torchmetrics.image import PeakSignalNoiseRatio, StructuralSimilarityIndexMeasure +from torchmetrics.image import PeakSignalNoiseRatio, StructuralSimilarityIndexMeasure,MultiScaleStructuralSimilarityIndexMeasure from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity from nerfstudio.cameras.cameras import Cameras from gsplat._torch_impl import quat_to_rotmat @@ -42,6 +42,7 @@ from nerfstudio.cameras.camera_optimizers import CameraOptimizer, CameraOptimize from torchmetrics.image import StructuralSimilarityIndexMeasure from gsplat.rasterize import RasterizeGaussians from gsplat.project_gaussians import ProjectGaussians +from nerfstudio.model_components.losses import depth_ranking_loss from gsplat.sh import SphericalHarmonics, num_sh_bases @@ -96,9 +97,9 @@ class GaussianSplattingModelConfig(ModelConfig): """training starts at 1/d resolution, every n steps this is doubled""" num_downscales: int = 2 """at the beginning, resolution is 1/2^d, where d is this number""" - cull_alpha_thresh: float = 0.05 + cull_alpha_thresh: float = 0.1 """threshold of opacity for culling gaussians""" - cull_scale_thresh: float = 0.1 + cull_scale_thresh: float = 0.5 """threshold of scale for culling gaussians""" reset_alpha_every: int = 30 """Every this many refinement steps, reset the alpha""" @@ -120,7 +121,7 @@ class GaussianSplattingModelConfig(ModelConfig): """number of extra points to add to the model in addition to sfm points, randomly distributed""" ssim_lambda: float = 0.2 """weight of ssim loss""" - stop_split_at: int = 25000 + stop_split_at: int = 15000 """stop splitting at this step""" sh_degree: int = 4 """maximum degree of spherical harmonics to use""" @@ -297,7 +298,9 @@ class GaussianSplattingModel(Model): def refinement_after(self, optimizers: Optimizers, step): if self.step >= self.config.warmup_length: with torch.no_grad(): - if self.step < self.config.stop_split_at: + # only split/cull if we've seen every image since opacity reset + reset_interval = self.config.reset_alpha_every * self.config.refine_every + if self.step < self.config.stop_split_at and self.step % reset_interval > self.num_train_data + self.config.refine_every: # then we densify avg_grad_norm = ( (self.xys_grad_norm / self.vis_counts) * 0.5 * max(self.last_size[0], self.last_size[1]) @@ -342,17 +345,17 @@ class GaussianSplattingModel(Model): for group, param in param_groups.items(): self.dup_in_optim(optimizers.optimizers[group], dup_idcs, param, 1) - # only cull if we've seen every image since opacity reset - reset_interval = self.config.reset_alpha_every * self.config.refine_every - if self.step % reset_interval > self.num_train_data: + # Offset all the opacity reset logic by refine_every so that we don't + # save checkpoints right when the opacity is reset (saves every 2k) + if self.step % reset_interval > self.num_train_data + self.config.refine_every: # then cull deleted_mask = self.cull_gaussians() param_groups = self.get_gaussian_param_groups() for group, param in param_groups.items(): self.remove_from_optim(optimizers.optimizers[group], deleted_mask, param) - if (self.step - 1) // self.config.refine_every % self.config.reset_alpha_every == 0: - reset_value = 0.01 + if self.step % reset_interval == self.config.refine_every: + reset_value = self.config.cull_alpha_thresh * 0.8 self.opacities.data = torch.full_like( self.opacities.data, torch.logit(torch.tensor(reset_value)).item() ) @@ -365,11 +368,6 @@ class GaussianSplattingModel(Model): self.xys_grad_norm = None self.vis_counts = None self.max_2Dsize = None - # set gauss_scaling to the min size of each gaussian - #first convert the quats to rotation matrices to rotate the scales into world frame - # rots = quat_to_rotmat(self.quats.detach()/self.quats.detach().norm(dim=-1,keepdim=True)) - # scales_rot = torch.bmm(rots, torch.exp(self.scales.detach())[..., None]).squeeze() - # optimizers.optimizers['xyz'].gauss_scaling = self.scales.detach().exp().abs().amin(dim=-1,keepdim=True) def cull_gaussians(self): """ @@ -491,7 +489,7 @@ class GaussianSplattingModel(Model): def _get_downscale_factor(self): if self.training: return 2 ** max((self.config.num_downscales - self.step // self.config.resolution_schedule), 0) - else: + else: return 1 def get_outputs(self, camera: Cameras) -> Dict[str, Union[torch.Tensor, List]]: @@ -508,16 +506,17 @@ class GaussianSplattingModel(Model): print("Called get_outputs with not a camera") return {} assert camera.shape[0] == 1, "Only one camera at a time" - # dont mutate the input - camera = deepcopy(camera) - d = self._get_downscale_factor() - camera.rescale_output_resolution(1 / d) + if self.training: + #currently relies on the branch vickie/camera-grads + self.camera_optimizer.apply_to_camera(camera) if self.crop_box is not None and not self.training: crop_ids = self.crop_box.within(self.means).squeeze() if crop_ids.sum() == 0: return {"rgb": torch.full((camera.height.item(), camera.width.item(), 3), 0.5, device=self.device)} else: crop_ids = None + camera_downscale = self._get_downscale_factor() + camera.rescale_output_resolution(1 / camera_downscale) # shift the camera to center of scene looking at center R = camera.camera_to_worlds[0, :3, :3] # 3 x 3 T = camera.camera_to_worlds[0, :3, 3:4] # 3 x 1 @@ -577,6 +576,7 @@ class GaussianSplattingModel(Model): W, tile_bounds, ) + # if self.training: # # experimental gradient scaling for floater removal # depths_de = depths.detach() @@ -608,19 +608,20 @@ class GaussianSplattingModel(Model): background, ) depth_im = None - if not self.training: - depth_im = RasterizeGaussians.apply( - self.xys, - depths, - self.radii, - conics, - num_tiles_hit, - depths[:, None].repeat(1, 3), - torch.sigmoid(opacities_crop), - H, - W, - torch.ones(3, device=self.device) * 10, - )[..., 0:1] + depth_im = RasterizeGaussians.apply( + self.xys.detach(), + depths, + self.radii, + conics.detach(), + num_tiles_hit, + depths[:, None].repeat(1, 3), + torch.sigmoid(opacities_crop.detach()), + H, + W, + torch.ones(3, device=self.device) * 10, + )[..., 0:1] + # rescale the camera back to original dimensions + camera.rescale_output_resolution(camera_downscale) return {"rgb": rgb, "depth": depth_im} def get_metrics_dict(self, outputs, batch) -> Dict[str, torch.Tensor]: @@ -646,6 +647,7 @@ class GaussianSplattingModel(Model): metrics_dict["psnr"] = self.psnr(predicted_rgb, gt_rgb) self.camera_optimizer.get_metrics_dict(metrics_dict) + metrics_dict['gaussian_count'] = self.num_points return metrics_dict def get_loss_dict(self, outputs, batch, metrics_dict=None) -> Dict[str, torch.Tensor]: diff --git a/nerfstudio/pipelines/base_pipeline.py b/nerfstudio/pipelines/base_pipeline.py index fa3969f03..0f9388568 100644 --- a/nerfstudio/pipelines/base_pipeline.py +++ b/nerfstudio/pipelines/base_pipeline.py @@ -41,6 +41,7 @@ from nerfstudio.data.datamanagers.base_datamanager import ( VanillaDataManager, ) from nerfstudio.data.datamanagers.parallel_datamanager import ParallelDataManager +from nerfstudio.data.datamanagers.full_images_datamanager import FullImageDatamanager from nerfstudio.engine.callbacks import TrainingCallback, TrainingCallbackAttributes from nerfstudio.models.base_model import Model, ModelConfig from nerfstudio.utils import profiler @@ -372,7 +373,7 @@ class VanillaPipeline(Pipeline): """ self.eval() metrics_dict_list = [] - assert isinstance(self.datamanager, (VanillaDataManager, ParallelDataManager)) + assert isinstance(self.datamanager, (VanillaDataManager, ParallelDataManager,FullImageDatamanager)) num_images = len(self.datamanager.fixed_indices_eval_dataloader) with Progress( TextColumn("[progress.description]{task.description}"), diff --git a/nerfstudio/viewer_beta/viewer.py b/nerfstudio/viewer_beta/viewer.py index e48adf987..b063e1954 100644 --- a/nerfstudio/viewer_beta/viewer.py +++ b/nerfstudio/viewer_beta/viewer.py @@ -235,6 +235,7 @@ class Viewer: self.camera_handles[idx].visible = visible def update_camera_poses(self): + #TODO this fn accounts for like ~5% of total train time # Update the train camera locations based on optimization assert self.camera_handles is not None if hasattr(self.pipeline.datamanager, "train_camera_optimizer"): @@ -366,7 +367,6 @@ class Viewer: # this stops training while moving to make the response smoother while time.time() - self.last_move_time < 0.1: time.sleep(0.05) - # self.render_statemachine.action(RenderAction("static", self.camera_state)) if self.trainer is not None and self.trainer.training_state == "training" and self.train_util != 1: if ( EventName.TRAIN_RAYS_PER_SEC.value in GLOBAL_BUFFER["events"]