{"id":19544169,"url":"https://github.com/sunset1995/torch_efficient_distloss","last_synced_at":"2025-10-12T17:03:30.511Z","repository":{"id":38366056,"uuid":"502046404","full_name":"sunset1995/torch_efficient_distloss","owner":"sunset1995","description":"Efficient distortion loss with O(n) realization.","archived":false,"fork":false,"pushed_at":"2022-07-12T13:09:00.000Z","size":23,"stargazers_count":132,"open_issues_count":8,"forks_count":3,"subscribers_count":4,"default_branch":"main","last_synced_at":"2025-04-26T17:47:33.810Z","etag":null,"topics":["distortion-loss","efficient","mip-nerf-360","nerf"],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/sunset1995.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null}},"created_at":"2022-06-10T13:04:42.000Z","updated_at":"2025-04-18T17:54:42.000Z","dependencies_parsed_at":"2022-07-12T19:30:33.922Z","dependency_job_id":null,"html_url":"https://github.com/sunset1995/torch_efficient_distloss","commit_stats":null,"previous_names":[],"tags_count":2,"template":false,"template_full_name":null,"purl":"pkg:github/sunset1995/torch_efficient_distloss","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/sunset1995%2Ftorch_efficient_distloss","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/sunset1995%2Ftorch_efficient_distloss/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/sunset1995%2Ftorch_efficient_distloss/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/sunset1995%2Ftorch_efficient_distloss/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/sunset1995","download_url":"https://codeload.github.com/sunset1995/torch_efficient_distloss/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/sunset1995%2Ftorch_efficient_distloss/sbom","scorecard":null,"host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":272314765,"owners_count":24912463,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","status":"online","status_checked_at":"2025-08-27T02:00:09.397Z","response_time":76,"last_error":null,"robots_txt_status":"success","robots_txt_updated_at":"2025-07-24T06:49:26.215Z","robots_txt_url":"https://github.com/robots.txt","online":true,"can_crawl_api":true,"host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["distortion-loss","efficient","mip-nerf-360","nerf"],"created_at":"2024-11-11T03:24:58.103Z","updated_at":"2025-10-12T17:03:30.399Z","avatar_url":"https://github.com/sunset1995.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"# torch_efficient_distloss\n\nDistortion loss is proposed by [mip-nerf-360](https://jonbarron.info/mipnerf360/), which encourages volume rendering weights to be compact and sparse and can alleviate *floater* and *background collapse* artifact.\nIn our [DVGOv2 report](https://arxiv.org/abs/2206.05085), we show that the distortion loss is also helpful to point-based query, which speeds up our training and gives us better quantitative results.\n\nA pytorch pseudo-code for the distortion loss:\n```python\ndef original_distloss(w, m, interval):\n    '''\n    Original O(N^2) realization of distortion loss.\n    There are B rays each with N sampled points.\n    w:        Float tensor in shape [B,N]. Volume rendering weights of each point.\n    m:        Float tensor in shape [B,N]. Midpoint distance to camera of each point.\n    interval: Scalar or float tensor in shape [B,N]. The query interval of each point.\n    '''\n    loss_uni = (1/3) * (interval * w.pow(2)).sum(-1).mean()\n    ww = w.unsqueeze(-1) * w.unsqueeze(-2)          # [B,N,N]\n    mm = (m.unsqueeze(-1) - m.unsqueeze(-2)).abs()  # [B,N,N]\n    loss_bi = (ww * mm).sum((-1,-2)).mean()\n    return loss_uni + loss_bi\n```\n\nUnfortunately, 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.\n\nPlease cite mip-nerf-360 if you find this repo helpful. We will be happy if you also cite DVGOv2.\n```\n@inproceedings{BarronMVSH22,\n  author    = {Jonathan T. Barron and\n               Ben Mildenhall and\n               Dor Verbin and\n               Pratul P. Srinivasan and\n               Peter Hedman},\n  title     = {Mip-NeRF 360: Unbounded Anti-Aliased Neural Radiance Fields},\n  booktitle = {CVPR},\n  year      = {2022},\n}\n\n@article{SunSC22_2,\n  author    = {Cheng Sun and\n               Min Sun and\n               Hwann{-}Tzong Chen},\n  title     = {Improved Direct Voxel Grid Optimization for Radiance Fields Reconstruction},\n  journal   = {arxiv cs.GR 2206.05085},\n  year      = {2022},\n}\n```\n\n## Installation\n```\npip install torch_efficient_distloss\n```\nAssumed `Pytorch` and `numpy` are already installed.\n\n## Documentation\nAll functions are runs in `O(N)` and are numerical equivalent to the distortion loss.\n```python\nimport torch\nfrom torch_efficient_distloss import eff_distloss, eff_distloss_native, flatten_eff_distloss\n\n# A toy example\nB = 8192  # number of rays\nN = 128   # number of points on a ray\nw = torch.rand(B, N).cuda()\nw = w / w.sum(-1, keepdim=True)\nw = w.clone().requires_grad_()\ns = torch.linspace(0, 1, N+1).cuda()\nm = (s[1:] + s[:-1]) * 0.5\nm = m[None].repeat(B,1)\ninterval = 1/N\n\nloss = 0.01 * eff_distloss(w, m, interval)\nloss.backward()\nprint('Loss', loss)\nprint('Gradient', w.grad)\n```\n- `eff_distloss_native`:\n    - Using built-in Pytorch operation to implement the `O(N)` distortion loss.\n    - Input:\n        - `w`: Float tensor in shape [B,N]. Volume rendering weights of each point.\n        - `m`: Float tensor in shape [B,N]. Midpoint distance to camera of each point.\n        - `interval`: Scalar or float tensor in shape [B,N]. The query interval of each point.\n- `eff_distloss`:\n    - The same as `eff_distloss_native`. Slightly faster and consume slightly more GPU memory.\n- `flatten_eff_distloss`:\n    - Support varied number of sampled points on each ray.\n    - All input tensor should be flatten.\n    - 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]\u003c=ray_id[i]`) and ranging from `0` to `N-1`.\n- Loss weight around `0.01` to `0.001` is recommended.\n\n\n## Testing\n### Numerical equivalent\nRun `python test.py`. All our implementation is numerical equivalent to the `O(N^2)` `original_distloss`.\n\n### Speed and memeory benchmark\nRun `python test_time_mem.py`. We use a batch of `B=8192` rays. Below is the results on my `RTX 2080Ti` GPU.\n- Peak GPU memory (MB)\n    | \\# of pts `N` | 32 | 64 | 128 | 256 | 384 | 512 |\n    |:------------:|:--:|:--:|:---:|:---:|:---:|:---:|\n    |`original_distloss`   |102|396|1560|6192|OOM|OOM|\n    |`eff_distloss_native` |12|24|48|96|144|192|\n    |`eff_distloss`        |14|28|56|112|168|224|\n    |`flatten_eff_distloss`|13|26|52|104|156|208|\n- Run time accumulated over 100 runs (sec)\n    | \\# of pts `N` | 32 | 64 | 128 | 256 | 384 | 512 |\n    |:------------:|:--:|:--:|:---:|:---:|:---:|:---:|\n    |`original_distloss`   |0.2|0.8|2.4|17.9|OOM|OOM|\n    |`eff_distloss_native` |0.1|0.1|0.1|0.2|0.3|0.3|\n    |`eff_distloss`        |0.1|0.1|0.1|0.1|0.2|0.2|\n    |`flatten_eff_distloss`|0.1|0.1|0.1|0.2|0.2|0.3|\n\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fsunset1995%2Ftorch_efficient_distloss","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fsunset1995%2Ftorch_efficient_distloss","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fsunset1995%2Ftorch_efficient_distloss/lists"}