misc fix, manually merge ashawkey#83, add basis based dnerf

assets/update_logs.md

@@ -1,5 +1,6 @@
 ## Update logs
 
+* 7.16: add temporal basis based dynamic nerf (experimental). It trains much faster compared to the deformation based dynamic nerf, but performance is much worse for now...
 * 6.29: add support for HyperNeRF's dataset.
  * we use a simplified pinhole camera model, may introduce bias.
 * 6.26: add support for D-NeRF.

dnerf/network_basis.py

@@ -0,0 +1,262 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from encoding import get_encoder
+from activation import trunc_exp
+from .renderer import NeRFRenderer
+
+
+class NeRFNetwork(NeRFRenderer):
+ def __init__(self,
+ encoding="tiledgrid",
+ encoding_dir="sphere_harmonics",
+ encoding_time="frequency",
+ encoding_bg="hashgrid",
+ num_layers=2,
+ hidden_dim=64,
+ geo_feat_dim=32,
+ num_layers_color=3,
+ hidden_dim_color=64,
+ num_layers_bg=2,
+ hidden_dim_bg=64,
+ sigma_basis_dim=32,
+ color_basis_dim=8,
+ num_layers_basis=5,
+ hidden_dim_basis=128,
+ bound=1,
+ **kwargs,
+ ):
+ super().__init__(bound, **kwargs)
+
+ # basis network
+ self.num_layers_basis = num_layers_basis
+ self.hidden_dim_basis = hidden_dim_basis
+ self.sigma_basis_dim = sigma_basis_dim
+ self.color_basis_dim = color_basis_dim
+ self.encoder_time, self.in_dim_time = get_encoder(encoding_time, input_dim=1, multires=6)
+
+ basis_net = []
+ for l in range(num_layers_basis):
+ if l == 0:
+ in_dim = self.in_dim_time
+ else:
+ in_dim = hidden_dim_basis
+
+ if l == num_layers_basis - 1:
+ out_dim = self.sigma_basis_dim + self.color_basis_dim
+ else:
+ out_dim = hidden_dim_basis
+
+ basis_net.append(nn.Linear(in_dim, out_dim, bias=False))
+
+ self.basis_net = nn.ModuleList(basis_net)
+
+ # sigma network
+ self.num_layers = num_layers
+ self.hidden_dim = hidden_dim
+ self.geo_feat_dim = geo_feat_dim
+ self.encoder, self.in_dim = get_encoder(encoding, desired_resolution=2048 * bound)
+
+ sigma_net = []
+ for l in range(num_layers):
+ if l == 0:
+ in_dim = self.in_dim
+ else:
+ in_dim = hidden_dim
+
+ if l == num_layers - 1:
+ out_dim = self.sigma_basis_dim + self.geo_feat_dim # SB sigma + features for color
+ else:
+ out_dim = hidden_dim
+
+ sigma_net.append(nn.Linear(in_dim, out_dim, bias=False))
+
+ self.sigma_net = nn.ModuleList(sigma_net)
+
+ # color network
+ self.num_layers_color = num_layers_color 
+ self.hidden_dim_color = hidden_dim_color
+ self.encoder_dir, self.in_dim_dir = get_encoder(encoding_dir)
+
+ color_net = []
+ for l in range(num_layers_color):
+ if l == 0:
+ in_dim = self.in_dim_dir + self.geo_feat_dim
+ else:
+ in_dim = hidden_dim
+
+ if l == num_layers_color - 1:
+ out_dim = 3 * self.color_basis_dim # 3 * CB rgb
+ else:
+ out_dim = hidden_dim
+
+ color_net.append(nn.Linear(in_dim, out_dim, bias=False))
+
+ self.color_net = nn.ModuleList(color_net)
+
+ # background network
+ if self.bg_radius > 0:
+ self.num_layers_bg = num_layers_bg 
+ self.hidden_dim_bg = hidden_dim_bg
+ self.encoder_bg, self.in_dim_bg = get_encoder(encoding_bg, input_dim=2, num_levels=4, log2_hashmap_size=19, desired_resolution=2048) # much smaller hashgrid 
+
+ bg_net = []
+ for l in range(num_layers_bg):
+ if l == 0:
+ in_dim = self.in_dim_bg + self.in_dim_dir
+ else:
+ in_dim = hidden_dim_bg
+
+ if l == num_layers_bg - 1:
+ out_dim = 3 # 3 rgb
+ else:
+ out_dim = hidden_dim_bg
+
+ bg_net.append(nn.Linear(in_dim, out_dim, bias=False))
+
+ self.bg_net = nn.ModuleList(bg_net)
+ else:
+ self.bg_net = None
+
+
+ def forward(self, x, d, t):
+ # x: [N, 3], in [-bound, bound]
+ # d: [N, 3], nomalized in [-1, 1]
+ # t: [1, 1], in [0, 1]
+
+ # time --> basis
+ enc_t = self.encoder_time(t) # [1, 1] --> [1, C']
+ h = enc_t
+ for l in range(self.num_layers_basis):
+ h = self.basis_net[l](h)
+ if l != self.num_layers_basis - 1:
+ h = F.relu(h, inplace=True)
+
+ sigma_basis = h[0, :self.sigma_basis_dim]
+ color_basis = h[0, self.sigma_basis_dim:]
+
+ # sigma
+ x = self.encoder(x, bound=self.bound)
+ h = x
+ for l in range(self.num_layers):
+ h = self.sigma_net[l](h)
+ if l != self.num_layers - 1:
+ h = F.relu(h, inplace=True)
+
+ sigma = trunc_exp(h[..., :self.sigma_basis_dim] @ sigma_basis)
+ geo_feat = h[..., self.sigma_basis_dim:]
+
+ # color
+ d = self.encoder_dir(d)
+ h = torch.cat([d, geo_feat], dim=-1)
+ for l in range(self.num_layers_color):
+ h = self.color_net[l](h)
+ if l != self.num_layers_color - 1:
+ h = F.relu(h, inplace=True)
+
+ # sigmoid activation for rgb
+ rgbs = torch.sigmoid(h.view(-1, 3, self.color_basis_dim) @ color_basis)
+
+ return sigma, rgbs, None
+
+ def density(self, x, t):
+ # x: [N, 3], in [-bound, bound]
+ # t: [1, 1], in [0, 1]
+
+ results = {}
+
+ # time --> basis
+ enc_t = self.encoder_time(t) # [1, 1] --> [1, C']
+ h = enc_t
+ for l in range(self.num_layers_basis):
+ h = self.basis_net[l](h)
+ if l != self.num_layers_basis - 1:
+ h = F.relu(h, inplace=True)
+
+ sigma_basis = h[0, :self.sigma_basis_dim]
+ color_basis = h[0, self.sigma_basis_dim:]
+
+ # sigma
+ x = self.encoder(x, bound=self.bound)
+ h = x
+ for l in range(self.num_layers):
+ h = self.sigma_net[l](h)
+ if l != self.num_layers - 1:
+ h = F.relu(h, inplace=True)
+
+ sigma = trunc_exp(h[..., :self.sigma_basis_dim] @ sigma_basis)
+ geo_feat = h[..., self.sigma_basis_dim:]
+
+ results['sigma'] = sigma
+ results['geo_feat'] = geo_feat
+ # results['color_basis'] = color_basis
+
+ return results
+
+ def background(self, x, d):
+ # x: [N, 2], in [-1, 1]
+
+ h = self.encoder_bg(x) # [N, C]
+ d = self.encoder_dir(d)
+
+ h = torch.cat([d, h], dim=-1)
+ for l in range(self.num_layers_bg):
+ h = self.bg_net[l](h)
+ if l != self.num_layers_bg - 1:
+ h = F.relu(h, inplace=True)
+
+ # sigmoid activation for rgb
+ rgbs = torch.sigmoid(h)
+
+ return rgbs
+
+ # TODO: non cuda-ray mode is broken for now... (how to pass color_basis to self.color())
+ # # allow masked inference
+ # def color(self, x, d, mask=None, geo_feat=None, **kwargs):
+ # # x: [N, 3] in [-bound, bound]
+ # # t: [1, 1], in [0, 1]
+ # # mask: [N,], bool, indicates where we actually needs to compute rgb.
+
+ # if mask is not None:
+ # rgbs = torch.zeros(mask.shape[0], 3, dtype=x.dtype, device=x.device) # [N, 3]
+ # # in case of empty mask
+ # if not mask.any():
+ # return rgbs
+ # x = x[mask]
+ # d = d[mask]
+ # geo_feat = geo_feat[mask]
+
+ # d = self.encoder_dir(d)
+ # h = torch.cat([d, geo_feat], dim=-1)
+ # for l in range(self.num_layers_color):
+ # h = self.color_net[l](h)
+ # if l != self.num_layers_color - 1:
+ # h = F.relu(h, inplace=True)
+
+ # # sigmoid activation for rgb
+ # h = torch.sigmoid(h)
+
+ # if mask is not None:
+ # rgbs[mask] = h.to(rgbs.dtype) # fp16 --> fp32
+ # else:
+ # rgbs = h
+
+ # return rgbs 
+
+ # optimizer utils
+ def get_params(self, lr, lr_net):
+
+ params = [
+ {'params': self.encoder.parameters(), 'lr': lr},
+ {'params': self.sigma_net.parameters(), 'lr': lr_net},
+ {'params': self.encoder_dir.parameters(), 'lr': lr},
+ {'params': self.color_net.parameters(), 'lr': lr_net},
+ {'params': self.encoder_time.parameters(), 'lr': lr},
+ {'params': self.basis_net.parameters(), 'lr': lr_net},
+ ]
+ if self.bg_radius > 0:
+ params.append({'params': self.encoder_bg.parameters(), 'lr': lr})
+ params.append({'params': self.bg_net.parameters(), 'lr': lr_net})
+
+ return params

dnerf/utils.py

@@ -115,8 +115,8 @@ def train_step(self, data):
  loss = loss.mean()
 
  # deform regularization
- deform = outputs['deform']
- loss = loss + 1e-3 * deform.abs().mean()
+ if 'deform' in outputs and outputs['deform'] is not None:
+  loss = loss + 1e-3 * outputs['deform'].abs().mean()
 
  return pred_rgb, gt_rgb, loss
 

main_dnerf.py

@@ -34,6 +34,7 @@
 
  ### network backbone options
  parser.add_argument('--fp16', action='store_true', help="use amp mixed precision training")
+ parser.add_argument('--basis', action='store_true', help="[experimental] use temporal basis instead of deformation to model dynamic scene (check Fourier PlenOctree and NeuVV)")
  # parser.add_argument('--ff', action='store_true', help="use fully-fused MLP")
  # parser.add_argument('--tcnn', action='store_true', help="use TCNN backend")
 
@@ -69,7 +70,11 @@
  opt.cuda_ray = True
  opt.preload = True
 
- from dnerf.network import NeRFNetwork
+ if opt.basis:
+ assert opt.cuda_ray, "Non-cuda-ray mode is temporarily broken with temporal basis mode"
+ from dnerf.network_basis import NeRFNetwork
+ else:
+ from dnerf.network import NeRFNetwork
 
  print(opt)
 

nerf/renderer.py

@@ -249,6 +249,7 @@ def run(self, rays_o, rays_d, num_steps=128, upsample_steps=128, bg_color=None,
  return {
  'depth': depth,
  'image': image,
+ 'weights_sum': weights_sum,
  }
 
 
@@ -274,6 +275,8 @@ def run_cuda(self, rays_o, rays_d, dt_gamma=0, bg_color=None, perturb=False, for
  elif bg_color is None:
  bg_color = 1
 
+ results = {}
+
  if self.training:
  # setup counter
  counter = self.step_counter[self.local_step % 16]
@@ -314,6 +317,8 @@ def run_cuda(self, rays_o, rays_d, dt_gamma=0, bg_color=None, perturb=False, for
  depth = torch.clamp(depth - nears, min=0) / (fars - nears)
  image = image.view(*prefix, 3)
  depth = depth.view(*prefix)
+
+ results['weights_sum'] = weights_sum
 
  else:
 
@@ -365,11 +370,11 @@ def run_cuda(self, rays_o, rays_d, dt_gamma=0, bg_color=None, perturb=False, for
  depth = torch.clamp(depth - nears, min=0) / (fars - nears)
  image = image.view(*prefix, 3)
  depth = depth.view(*prefix)
+
+ results['depth'] = depth
+ results['image'] = image
 
- return {
- 'depth': depth,
- 'image': image,
- }
+ return results
 
  @torch.no_grad()
  def mark_untrained_grid(self, poses, intrinsic, S=64):

nerf/utils.py

@@ -292,7 +292,7 @@ def __init__(self,
  self.scaler = torch.cuda.amp.GradScaler(enabled=self.fp16)
 
  # variable init
- self.epoch = 1
+ self.epoch = 0
  self.global_step = 0
  self.local_step = 0
  self.stats = {
@@ -442,6 +442,11 @@ def train_step(self, data):
 
  loss = loss.mean()
 
+ # extra loss
+ # pred_weights_sum = outputs['weights_sum'] + 1e-8
+ # loss_ws = - 1e-1 * pred_weights_sum * torch.log(pred_weights_sum) # entropy to encourage weights_sum to be 0 or 1.
+ # loss = loss + loss_ws.mean()
+
  return pred_rgb, gt_rgb, loss
 
  def eval_step(self, data):
@@ -523,7 +528,7 @@ def train(self, train_loader, valid_loader, max_epochs):
  # get a ref to error_map
  self.error_map = train_loader._data.error_map
 
- for epoch in range(self.epoch, max_epochs + 1):
+ for epoch in range(self.epoch + 1, max_epochs + 1):
  self.epoch = epoch
 
  self.train_one_epoch(train_loader)

readme.md

@@ -187,8 +187,12 @@ python main_CCNeRF.py data/nerf_synthetic/hotdog --workspace trial_cc_hotdog -O
 
 ### D-NeRF
 # almost the same as Instant-ngp NeRF, just replace the main script.
+# use deformation to model dynamic scene
 python main_dnerf.py data/dnerf/jumpingjacks --workspace trial_dnerf_jumpingjacks -O --bound 1.0 --scale 0.8 --dt_gamma 0
 python main_dnerf.py data/dnerf/jumpingjacks --workspace trial_dnerf_jumpingjacks -O --bound 1.0 --scale 0.8 --dt_gamma 0 --gui
+# use temporal basis to model dynamic scene
+python main_dnerf.py data/dnerf/jumpingjacks --workspace trial_dnerf_basis_jumpingjacks -O --bound 1.0 --scale 0.8 --dt_gamma 0 --basis
+python main_dnerf.py data/dnerf/jumpingjacks --workspace trial_dnerf_basis_jumpingjacks -O --bound 1.0 --scale 0.8 --dt_gamma 0 --basis --gui
 # for the hypernerf dataset, first convert it into nerf-compatible format:
 python scripts/hyper2nerf.py data/split-cookie --downscale 2 # will generate transforms*.json
 python main_dnerf.py data/split-cookie/ --workspace trial_dnerf_cookies -O --bound 1 --scale 0.3 --dt_gamma 0