torch_efficient_distloss

Distortion loss is proposed by mip-nerf-360, which encourages volume rendering weights to be compact and sparse and can alleviate floater and background collapse artifact. 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:

def original_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()
    return loss_uni + loss_bi

Unfortunately, the straightforward implementation results in O(N^2) space time complexity for N sampled points on a ray. In this package, we provide our O(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},
}

Installation

pip install torch_efficient_distloss

Assumed Pytorch and numpy are already installed.

Documentation

All functions are runs in O(N) and are numerical equivalent to the distortion loss.

import torch
from torch_efficient_distloss import eff_distloss, eff_distloss_native, flatten_eff_distloss

# A toy example
B = 8192  # number of rays
N = 128   # number of points on a ray
w = 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.5
m = m[None].repeat(B,1)
interval = 1/N

loss = eff_distloss(w, m, interval)
loss.backward()
print('Loss', loss)
print('Gradient', w.grad)

• eff_distloss_native:
  • Using built-in Pytorch operation to implement the O(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 as eff_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 tensor ray_id to specify the ray index of each point. ray_id should be increasing (i.e., ray_id[i-1]<=ray_id[i]) and ranging from 0 to N-1.

Testing

Numerical equivalent

Run python test.py. All our implementation is numerical equivalent to the O(N^2) original_distloss.

Speed and memeory benchmark

Run python test_time_mem.py. We use a batch of B=8192 rays. Below is the results on my RTX 2080Ti GPU.

• Peak GPU memory (MB)

  # of pts N	32	64	128	256	384	512
  original_distloss	102	396	1560	6192	OOM	OOM
  eff_distloss_native	12	24	48	96	144	192
  eff_distloss	14	28	56	112	168	224
  flatten_eff_distloss	13	26	52	104	156	208

• Run time accumulated over 100 runs (sec)

  # of pts N	32	64	128	256	384	512
  original_distloss	0.2	0.8	2.4	17.9	OOM	OOM
  eff_distloss_native	0.1	0.1	0.1	0.2	0.3	0.3
  eff_distloss	0.1	0.1	0.1	0.1	0.2	0.2
  flatten_eff_distloss	0.1	0.1	0.1	0.2	0.2	0.3