tink/nerfstudio
4
0 Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-13 19:33:46 +08:00
Code Activity

Pose optimization (#721)

Browse Source 
* pose opt default nerfacto.

* added SO3 cam pose optimization which is slightly faster (~2.2m/iter) than BARF

* integrating record3d

* fix linting errors

* fix save checkpoint logic

* added record3d data

* workaround to avoid flag

* moved dict

* black

* record3d docs

* re-add typing to exp_map.

* Pose optimization refactor

* Fix camera optimizer param group logic

* Update docs + comments

* Update docs, shorten names

* Appease pylint

* repeat -> tile, apply missing name change

* Fix optimizer as scheduler bug

* Fix test

Co-authored-by: akristoffersen <akristoffersen@berkeley.edu>
Co-authored-by: Justin Kerr <justin.g.kerr@gmail.com>
Co-authored-by: Brent Yi <yibrenth@gmail.com>
This commit is contained in:
Ethan WeberakristoffersenJustin KerrBrent Yi
 authored and GitHub committed 2022-10-07 10:11:33 -07:00
1 parent 2531242d1d
commit cac177717f
8 files changed
+230 -103
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

docs/developer_guides/pipelines/datamanagers.md
+1 -1
View File
@@ -64,7 +64,7 @@ class VanillaDataManagerConfig(InstantiateConfig):
"""number of images to sample during eval iteration"""
eval_image_indices: Optional[Tuple[int, ...]] = (0,)
"""specifies the image indices to use during eval; if None, uses all"""
train_camera_optimizer: CameraOptimizerConfig = CameraOptimizerConfig()
camera_optimizer: CameraOptimizerConfig = CameraOptimizerConfig()
"""specifies the camera pose optimizer used during training"""
```
docs/quickstart/custom_dataset.md
+15 -1
View File
@@ -119,5 +119,19 @@ Record a video and export with the EXR + JPG sequence format.
Then, move the exported capture folder from your iPhone to your computer.
```
ns-train nerfacto record3d-data --data {RECORD3D_CAPTURE_DIR/EXR_RGBD}
ns-train nerfacto --pipeline.datamanager.train-camera-optimizer.mode SO3xR3 record3d-data --data {RECORD3D_CAPTURE_DIR/EXR_RGBD}
```
You can try with our data to see the correct formatting.
```shell
ns-download-data --dataset record3d
```
This will download data to `data/record3d/bear`. Then you can train, and don't forget to open up the viewer.
```python
ns-train nerfacto --pipeline.datamanager.train-camera-optimizer.mode SO3xR3 record3d-data --data data/record3d/bear
```
Notice we are adding the option `--pipeline.datamanager.train-camera-optimizer.mode SO3xR3`, which is experimental but turns on camera pose optimization for higher quality results with record3d. This is because the camera poses from record3d are quite noisy, unlike those from COLMAP (from the `ns-process-data` script).
nerfstudio/cameras/camera_optimizers.py
+60 -80
View File
@@ -18,32 +18,47 @@ Pose and Intrinsics Optimizers
from __future__ import annotations
import functools
from dataclasses import dataclass, field
from typing import Optional, Type, Union
from typing import Literal, Type, Union
import torch
import tyro
from torch import nn
from torchtyping import TensorType
from typing_extensions import assert_never
from nerfstudio.cameras.lie_groups import exp_map_SE3, exp_map_SO3xR3
from nerfstudio.configs import base_config as cfg
from nerfstudio.engine.optimizers import AdamOptimizerConfig
from nerfstudio.engine.schedulers import SchedulerConfig
from nerfstudio.utils import poses as pose_utils
# Camera Optimization related configs
@dataclass
class CameraOptimizerConfig(cfg.InstantiateConfig):
"""Default camera optimizer config. Note: This is a no-op class and will not optimize cameras."""
"""Configuration of optimization for camera poses."""
_target: Type = field(default_factory=lambda: CameraOptimizer)
mode: Literal["off", "SO3xR3", "SE3"] = "off"
"""Pose optimization strategy to use. If enabled, we recommend SO3xR3."""
@dataclass
class BARFPoseOptimizerConfig(CameraOptimizerConfig):
"""BARF camera optimizer."""
position_noise_std: float = 0.0
"""Noise to add to initial positions. Useful for debugging."""
_target: Type = field(default_factory=lambda: BARFOptimizer)
noise_variance: Optional[float] = None
"""Optional additional noise to add to poses. Useful for debugging pose optimization."""
orientation_noise_std: float = 0.0
"""Noise to add to initial orientations. Useful for debugging."""
optimizer: AdamOptimizerConfig = AdamOptimizerConfig(lr=6e-4, eps=1e-15)
"""ADAM parameters for camera optimization."""
scheduler: SchedulerConfig = SchedulerConfig(max_steps=10000)
"""Learning rate scheduler for camera optimizer.."""
param_group: tyro.conf.Suppress[str] = "camera_opt"
"""Name of the parameter group used for pose optimization. Can be any string that doesn't conflict with other
groups."""
class CameraOptimizer(nn.Module):
@@ -63,6 +78,24 @@ class CameraOptimizer(nn.Module):
self.num_cameras = num_cameras
self.device = device
# Initialize learnable parameters.
if self.config.mode == "off":
pass
elif self.config.mode in ("SO3xR3", "SE3"):
self.pose_adjustment = torch.nn.Parameter(torch.zeros((num_cameras, 6), device=device))
else:
assert_never(self.config.mode)
# Initialize pose noise; useful for debugging.
if config.position_noise_std != 0.0 or config.orientation_noise_std != 0.0:
assert config.position_noise_std >= 0.0 and config.orientation_noise_std >= 0.0
std_vector = torch.tensor(
[config.position_noise_std] * 3 + [config.orientation_noise_std] * 3, device=device
)
self.pose_noise = exp_map_SE3(torch.normal(torch.zeros((num_cameras, 6), device=device), std_vector))
else:
self.pose_noise = None
def forward(
self,
indices: TensorType["num_cameras"],
@@ -74,77 +107,24 @@ class CameraOptimizer(nn.Module):
Tranformation matrices from optimized camera coordinates coordinates
to given camera coordinates.
"""
return torch.eye(4, device=self.device).repeat(indices.shape[0], 1, 1)[:, :3, :4] # no-op (Identity Transform)
outputs = []
# Apply learned transformation delta.
if self.config.mode == "off":
pass
elif self.config.mode == "SO3xR3":
outputs.append(exp_map_SO3xR3(self.pose_adjustment[indices, :]))
elif self.config.mode == "SE3":
outputs.append(exp_map_SE3(self.pose_adjustment[indices, :]))
else:
assert_never(self.config.mode)
class BARFOptimizer(CameraOptimizer):
"""Bundle-Adjusting NeRF (BARF) Pose Optimization"""
# Apply initial pose noise.
if self.pose_noise is not None:
outputs.append(self.pose_noise[indices, :, :])
config: BARFPoseOptimizerConfig
def __init__(self, config: BARFPoseOptimizerConfig, num_cameras: int, device: Union[torch.device, str]) -> None:
super().__init__(config, num_cameras, device)
if self.config.noise_variance is not None:
pose_noise = torch.normal(torch.zeros(self.num_cameras, 6), self.config.noise_variance)
self.pose_noise = self.exp_map(pose_noise).detach().to(device)
self.pose_adjustment = nn.Embedding(self.num_cameras, 6, device=device)
nn.init.zeros_(self.pose_adjustment.weight)
@classmethod
def exp_map(cls, tangent_vector: TensorType["num_cameras", 6]) -> TensorType["num_cameras", 3, 4]:
"""Convert SE3 vector into [R|t] transformation matrix.
Args:
tangent_vector: SE3 vector
Returns:
Respective [R|t] tranformation matrices.
"""
tangent_vector_lin = tangent_vector[:, :3].view(-1, 3, 1)
tangent_vector_ang = tangent_vector[:, 3:].view(-1, 3, 1)
theta = torch.linalg.norm(tangent_vector_ang, dim=1).unsqueeze(1)
theta2 = theta**2
theta3 = theta**3
near_zero = theta < 1e-2
non_zero = torch.ones(1, dtype=tangent_vector.dtype, device=tangent_vector.device)
theta_nz = torch.where(near_zero, non_zero, theta)
theta2_nz = torch.where(near_zero, non_zero, theta2)
theta3_nz = torch.where(near_zero, non_zero, theta3)
# Compute the rotation
sine = theta.sin()
cosine = torch.where(near_zero, 8 / (4 + theta2) - 1, theta.cos())
sine_by_theta = torch.where(near_zero, 0.5 * cosine + 0.5, sine / theta_nz)
one_minus_cosine_by_theta2 = torch.where(near_zero, 0.5 * sine_by_theta, (1 - cosine) / theta2_nz)
ret = torch.zeros(tangent_vector.shape[0], 3, 4).to(dtype=tangent_vector.dtype, device=tangent_vector.device)
ret[:, :3, :3] = one_minus_cosine_by_theta2 * tangent_vector_ang @ tangent_vector_ang.transpose(1, 2)
ret[:, 0, 0] += cosine.view(-1)
ret[:, 1, 1] += cosine.view(-1)
ret[:, 2, 2] += cosine.view(-1)
temp = sine_by_theta.view(-1, 1) * tangent_vector_ang.view(-1, 3)
ret[:, 0, 1] -= temp[:, 2]
ret[:, 1, 0] += temp[:, 2]
ret[:, 0, 2] += temp[:, 1]
ret[:, 2, 0] -= temp[:, 1]
ret[:, 1, 2] -= temp[:, 0]
ret[:, 2, 1] += temp[:, 0]
# Compute the translation
sine_by_theta = torch.where(near_zero, 1 - theta2 / 6, sine_by_theta)
one_minus_cosine_by_theta2 = torch.where(near_zero, 0.5 - theta2 / 24, one_minus_cosine_by_theta2)
theta_minus_sine_by_theta3_t = torch.where(near_zero, 1.0 / 6 - theta2 / 120, (theta - sine) / theta3_nz)
ret[:, :, 3:] = sine_by_theta * tangent_vector_lin
ret[:, :, 3:] += one_minus_cosine_by_theta2 * torch.cross(tangent_vector_ang, tangent_vector_lin, dim=1)
ret[:, :, 3:] += theta_minus_sine_by_theta3_t * (
tangent_vector_ang @ (tangent_vector_ang.transpose(1, 2) @ tangent_vector_lin)
)
return ret
def forward(self, indices: TensorType["num_cameras"]) -> TensorType["num_cameras", 3, 4]:
c_opt2c = self.exp_map(self.pose_adjustment.weight[indices])
if self.config.noise_variance is not None:
c_opt2c = pose_utils.multiply(c_opt2c, self.pose_noise[indices])
return c_opt2c
# Return: identity if no transforms are needed, otherwise multiply transforms together.
if len(outputs) == 0:
# Note that using repeat() instead of tile() here would result in unnecessary copies.
return torch.eye(4, device=self.device)[None, :3, :4].tile(indices.shape[0], 1, 1)
return functools.reduce(pose_utils.multiply, outputs)
nerfstudio/cameras/lie_groups.py
+115
View File
@@ -0,0 +1,115 @@
# Copyright 2022 The Nerfstudio Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Helper for Lie group operations. Currently only used for pose optimization.
"""
import torch
from torchtyping import TensorType
# We make an exception on snake case conventions because SO3 != so3.
def exp_map_SO3xR3(tangent_vector: TensorType["b", 6]) -> TensorType["b", 3, 4]: # pylint: disable=invalid-name
"""Compute the exponential map of the direct product group `SO(3) x R^3`.
This can be used for learning pose deltas on SE(3), and is generally faster than `exp_map_SE3`.
Args:
tangent_vector: Tangent vector; length-3 translations, followed by an `so(3)` tangent vector.
Returns:
[R|t] tranformation matrices.
"""
# code for SO3 map grabbed from pytorch3d and stripped down to bare-bones
log_rot = tangent_vector[:, 3:]
nrms = (log_rot * log_rot).sum(1)
rot_angles = torch.clamp(nrms, 1e-4).sqrt()
rot_angles_inv = 1.0 / rot_angles
fac1 = rot_angles_inv * rot_angles.sin()
fac2 = rot_angles_inv * rot_angles_inv * (1.0 - rot_angles.cos())
skews = torch.zeros((log_rot.shape[0], 3, 3), dtype=log_rot.dtype, device=log_rot.device)
skews[:, 0, 1] = -log_rot[:, 2]
skews[:, 0, 2] = log_rot[:, 1]
skews[:, 1, 0] = log_rot[:, 2]
skews[:, 1, 2] = -log_rot[:, 0]
skews[:, 2, 0] = -log_rot[:, 1]
skews[:, 2, 1] = log_rot[:, 0]
skews_square = torch.bmm(skews, skews)
ret = torch.zeros(tangent_vector.shape[0], 3, 4, dtype=tangent_vector.dtype, device=tangent_vector.device)
ret[:, :3, :3] = (
fac1[:, None, None] * skews
+ fac2[:, None, None] * skews_square
+ torch.eye(3, dtype=log_rot.dtype, device=log_rot.device)[None]
)
# Compute the translation
ret[:, :3, 3] = tangent_vector[:, :3]
return ret
def exp_map_SE3(tangent_vector: TensorType["b", 6]) -> TensorType["b", 3, 4]: # pylint: disable=invalid-name
"""Compute the exponential map `se(3) -> SE(3)`.
This can be used for learning pose deltas on `SE(3)`.
Args:
tangent_vector: A tangent vector from `se(3)`.
Returns:
[R|t] tranformation matrices.
"""
tangent_vector_lin = tangent_vector[:, :3].view(-1, 3, 1)
tangent_vector_ang = tangent_vector[:, 3:].view(-1, 3, 1)
theta = torch.linalg.norm(tangent_vector_ang, dim=1).unsqueeze(1)
theta2 = theta**2
theta3 = theta**3
near_zero = theta < 1e-2
non_zero = torch.ones(1, dtype=tangent_vector.dtype, device=tangent_vector.device)
theta_nz = torch.where(near_zero, non_zero, theta)
theta2_nz = torch.where(near_zero, non_zero, theta2)
theta3_nz = torch.where(near_zero, non_zero, theta3)
# Compute the rotation
sine = theta.sin()
cosine = torch.where(near_zero, 8 / (4 + theta2) - 1, theta.cos())
sine_by_theta = torch.where(near_zero, 0.5 * cosine + 0.5, sine / theta_nz)
one_minus_cosine_by_theta2 = torch.where(near_zero, 0.5 * sine_by_theta, (1 - cosine) / theta2_nz)
ret = torch.zeros(tangent_vector.shape[0], 3, 4).to(dtype=tangent_vector.dtype, device=tangent_vector.device)
ret[:, :3, :3] = one_minus_cosine_by_theta2 * tangent_vector_ang @ tangent_vector_ang.transpose(1, 2)
ret[:, 0, 0] += cosine.view(-1)
ret[:, 1, 1] += cosine.view(-1)
ret[:, 2, 2] += cosine.view(-1)
temp = sine_by_theta.view(-1, 1) * tangent_vector_ang.view(-1, 3)
ret[:, 0, 1] -= temp[:, 2]
ret[:, 1, 0] += temp[:, 2]
ret[:, 0, 2] += temp[:, 1]
ret[:, 2, 0] -= temp[:, 1]
ret[:, 1, 2] -= temp[:, 0]
ret[:, 2, 1] += temp[:, 0]
# Compute the translation
sine_by_theta = torch.where(near_zero, 1 - theta2 / 6, sine_by_theta)
one_minus_cosine_by_theta2 = torch.where(near_zero, 0.5 - theta2 / 24, one_minus_cosine_by_theta2)
theta_minus_sine_by_theta3_t = torch.where(near_zero, 1.0 / 6 - theta2 / 120, (theta - sine) / theta3_nz)
ret[:, :, 3:] = sine_by_theta * tangent_vector_lin
ret[:, :, 3:] += one_minus_cosine_by_theta2 * torch.cross(tangent_vector_ang, tangent_vector_lin, dim=1)
ret[:, :, 3:] += theta_minus_sine_by_theta3_t * (
tangent_vector_ang @ (tangent_vector_ang.transpose(1, 2) @ tangent_vector_lin)
)
return ret
nerfstudio/data/datamanagers.py
+17 -12
View File
@@ -235,21 +235,22 @@ class VanillaDataManagerConfig(InstantiateConfig):
"""
_target: Type = field(default_factory=lambda: VanillaDataManager)
"""target class to instantiate"""
"""Target class to instantiate."""
dataparser: AnnotatedDataParserUnion = BlenderDataParserConfig()
"""specifies the dataparser used to unpack the data"""
"""Specifies the dataparser used to unpack the data."""
train_num_rays_per_batch: int = 1024
"""number of rays per batch to use per training iteration"""
"""Number of rays per batch to use per training iteration."""
train_num_images_to_sample_from: int = -1
"""number of images to sample during training iteration"""
"""Number of images to sample during training iteration."""
eval_num_rays_per_batch: int = 1024
"""number of rays per batch to use per eval iteration"""
"""Number of rays per batch to use per eval iteration."""
eval_num_images_to_sample_from: int = -1
"""number of images to sample during eval iteration"""
"""Number of images to sample during eval iteration."""
eval_image_indices: Optional[Tuple[int, ...]] = (0,)
"""specifies the image indices to use during eval; if None, uses all"""
train_camera_optimizer: CameraOptimizerConfig = CameraOptimizerConfig()
"""specifies the camera pose optimizer used during training"""
"""Specifies the image indices to use during eval; if None, uses all."""
camera_optimizer: CameraOptimizerConfig = CameraOptimizerConfig()
"""Specifies the camera pose optimizer used during training. Helpful if poses are noisy, such as for data from
Record3D."""
class VanillaDataManager(DataManager): # pylint: disable=abstract-method
@@ -309,7 +310,7 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method
)
self.iter_train_image_dataloader = iter(self.train_image_dataloader)
self.train_pixel_sampler = PixelSampler(self.config.train_num_rays_per_batch)
self.train_camera_optimizer = self.config.train_camera_optimizer.setup(
self.train_camera_optimizer = self.config.camera_optimizer.setup(
num_cameras=self.train_dataset.dataparser_outputs.cameras.size, device=self.device
)
self.train_ray_generator = RayGenerator(
@@ -389,6 +390,10 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method
param_groups = {}
camera_opt_params = list(self.train_camera_optimizer.parameters())
if len(camera_opt_params) > 0:
param_groups["camera_opt"] = list(camera_opt_params)
if self.config.camera_optimizer.mode != "off":
assert len(camera_opt_params) > 0
param_groups[self.config.camera_optimizer.param_group] = camera_opt_params
else:
assert len(camera_opt_params) == 0
return param_groups
nerfstudio/engine/optimizers.py
+15 -5
View File
@@ -24,13 +24,13 @@ import torch
from torch.cuda.amp.grad_scaler import GradScaler
from torch.nn.parameter import Parameter
from nerfstudio.configs.base_config import PrintableConfig
from nerfstudio.configs import base_config
from nerfstudio.utils import writer
# Optimizer related configs
@dataclass
class OptimizerConfig(PrintableConfig):
class OptimizerConfig(base_config.PrintableConfig):
"""Basic optimizer config with RAdam"""
_target: Type = torch.optim.Adam
@@ -60,17 +60,27 @@ class RAdamOptimizerConfig(OptimizerConfig):
_target: Type = torch.optim.RAdam
def setup_optimizers(config: Dict[str, Any], param_groups: Dict[str, List[Parameter]]) -> "Optimizers":
def setup_optimizers(config: base_config.Config, param_groups: Dict[str, List[Parameter]]) -> "Optimizers":
"""Helper to set up the optimizers
Args:
config: The optimizer configuration object.
config: The trainer configuration object.
param_groups: A dictionary of parameter groups to optimize.
Returns:
The optimizers object.
"""
return Optimizers(config, param_groups)
optimizer_config = config.optimizers.copy()
# Add the camera optimizer if enabled.
camera_optimizer_config = config.pipeline.datamanager.camera_optimizer
if camera_optimizer_config.mode != "off":
assert camera_optimizer_config.param_group not in optimizer_config
optimizer_config[camera_optimizer_config.param_group] = {
"optimizer": config.pipeline.datamanager.camera_optimizer.optimizer,
"scheduler": config.pipeline.datamanager.camera_optimizer.scheduler,
}
return Optimizers(optimizer_config, param_groups)
class Optimizers:
nerfstudio/engine/trainer.py
+1 -1
View File
@@ -129,7 +129,7 @@ class Trainer:
self.pipeline: VanillaPipeline = self.config.pipeline.setup(
device=self.device, test_mode=test_mode, world_size=self.world_size, local_rank=self.local_rank
)
self.optimizers = setup_optimizers(self.config.optimizers, self.pipeline.get_param_groups())
self.optimizers = setup_optimizers(self.config, self.pipeline.get_param_groups())
self._load_checkpoint()
nerfstudio/pipelines/base_pipeline.py
+6 -3
View File
@@ -249,14 +249,17 @@ class VanillaPipeline(Pipeline):
ray_bundle, batch = self.datamanager.next_train(step)
model_outputs = self.model(ray_bundle)
metrics_dict = self.model.get_metrics_dict(model_outputs, batch)
if "camera_opt" in self.datamanager.get_param_groups():
camera_opt_param_group = self.config.datamanager.camera_optimizer.param_group
if camera_opt_param_group in self.datamanager.get_param_groups():
# Report the camera optimization metrics
metrics_dict["camera_opt_translation"] = (
self.datamanager.get_param_groups()["camera_opt"][0].data[:, :3].norm()
self.datamanager.get_param_groups()[camera_opt_param_group][0].data[:, :3].norm()
)
metrics_dict["camera_opt_rotation"] = (
self.datamanager.get_param_groups()["camera_opt"][0].data[:, 3:].norm()
self.datamanager.get_param_groups()[camera_opt_param_group][0].data[:, 3:].norm()
)
loss_dict = self.model.get_loss_dict(model_outputs, batch, metrics_dict)
return model_outputs, loss_dict, metrics_dict