Package: torch-efficient-distloss
torch_efficient_distlossDistortion loss is proposed bymip-nerf-360, which encourages volume rendering weights to be compact and sparse and can alleviatefloaterandbackground collapseartifact.
In our DVGOv2 report (coming soon), we show that the distortion loss is also helpful to point-based query, which speeds up our training and gives us better quantitative results.A pytorch pseudo-code for the distortion loss:deforiginal_distloss(w,m,interval):'''Original O(N^2) realization of distortion loss.There are B rays each with N sampled points.w:        Float tensor in shape [B,N]. Volume rendering weights of each point.m:        Float tensor in shape [B,N]. Midpoint distance to camera of each point.interval: Scalar or float tensor in shape [B,N]. The query interval of each point.'''loss_uni=(1/3)*(interval*w.pow(2)).sum(-1).mean()ww=w.unsqueeze(-1)*w.unsqueeze(-2)# [B,N,N]mm=(m.unsqueeze(-1)-m.unsqueeze(-2)).abs()# [B,N,N]loss_bi=(ww*mm).sum((-1,-2)).mean()returnloss_uni+loss_biUnfortunately, the straightforward implementation results inO(N^2)space time complexity for N sampled points on a ray. In this package, we provide ourO(N)realization presnted in the DVGOv2 report.Please cite mip-nerf-360 if you find this repo helpful. We will be happy if you also cite DVGOv2.@inproceedings{BarronMVSH22,
  author    = {Jonathan T. Barron and
               Ben Mildenhall and
               Dor Verbin and
               Pratul P. Srinivasan and
               Peter Hedman},
  title     = {Mip-NeRF 360: Unbounded Anti-Aliased Neural Radiance Fields},
  booktitle = {CVPR},
  year      = {2022},
}

@article{SunSC22_2,
  author    = {Cheng Sun and
               Min Sun and
               Hwann{-}Tzong Chen},
  title     = {Improved Direct Voxel Grid Optimization for Radiance Fields Reconstruction},
  journal   = {to be announced},
  year      = {2022},
}Installationpip install torch_efficient_distlossAssumedPytorchandnumpyare already installed.DocumentationAll functions are runs inO(N)and are numerical equivalent to the distortion loss.importtorchfromtorch_efficient_distlossimporteff_distloss,eff_distloss_native,flatten_eff_distloss# A toy exampleB=8192# number of raysN=128# number of points on a rayw=torch.rand(B,N).cuda()w=w/w.sum(-1,keepdim=True)w=w.clone().requires_grad_()s=torch.linspace(0,1,N+1).cuda()m=(s[1:]+s[:-1])*0.5m=m[None].repeat(B,1)interval=1/Nloss=eff_distloss(w,m,interval)loss.backward()print('Loss',loss)print('Gradient',w.grad)eff_distloss_native:Using built-in Pytorch operation to implement theO(N)distortion loss.Input:w: Float tensor in shape [B,N]. Volume rendering weights of each point.m: Float tensor in shape [B,N]. Midpoint distance to camera of each point.interval: Scalar or float tensor in shape [B,N]. The query interval of each point.eff_distloss:The same aseff_distloss_native. Slightly faster and consume slightly more GPU memory.flatten_eff_distloss:Support varied number of sampled points on each ray.All input tensor should be flatten.Should provide an additional flatten Long tensorray_idto specify the ray index of each point.ray_idshould be increasing (i.e.,ray_id[i-1]<=ray_id[i]) and ranging from0toN-1.TestingNumerical equivalentRunpython test.py. All our implementation is numerical equivalent to theO(N^2)original_distloss.Speed and memeory benchmarkRunpython test_time_mem.py. We use a batch ofB=8192rays. Below is the results on myRTX 2080TiGPU.Peak GPU memory (MB)# of ptsN3264128256384512original_distloss10239615606192OOMOOMeff_distloss_native12244896144192eff_distloss142856112168224flatten_eff_distloss132652104156208Run time accumulated over 100 runs (sec)# of ptsN3264128256384512original_distloss0.20.82.417.9OOMOOMeff_distloss_native0.10.10.10.20.30.3eff_distloss0.10.10.10.10.20.2flatten_eff_distloss0.10.10.10.20.20.3