ending2015a/hash-grid-encoding

Pure PyTorch implementation of Nvidia's hash grid encoding: https://nvlabs.github.io/instant-ngp/

Hash Grid Encoding

This repo contains an implementation of NVidia's hash grid encoding from Instant Neural Graphics Primitives and a runnable example of gigapixel tasks. The hash grid is implemented in pure PyTorch hense it's more human friendly than NVidia's original implementation (C++/CUDA).

Some features:

• Implemented in pure PyTorch.
• Supports arbitrary dimensions.

How To Use

MultiResHashGrid

To use the MultiResHashGrid in your own project, you can simply copy-paste the code in encoding.py into your project. For example:

import torch
import encoding

enc = encoding.MultiResHashGrid(2)  # 2D image data
enc = encoding.MultiResHashGrid(3)  # 3D data

dim = 3
batch_size = 100

# The input value must be within the range [0, 1]
input = torch.rand((batch_size, dim), dtpye=torch.float32)
enc_input = enc(input)

# Then you can forward into your network
model = MyMLP(dim=enc_input.output_dim, out_dim=1)
output = model(enc_input)

# Move to other devices
enc = enc.to(dtype='cuda')

Gigapixel task

This repo also contains a runnable gigapixel image task, which is implemented based on PyTorch Lightning. For more instructions of running this code, see Examples.

Examples

Albert

Run this example:

python train.py -i data/albert.jpg --enc_method hashgrid --visualize

Tokyo

tokyo-zoom.mp4

Download the tokyo image and place it at data/tokyo.jpg.

To run the tokyo example in its original size (56718 x 21450 pixels), your GPU must have memory at least 20GB. If your GPU have no such amount of memory, you can use the convert.py script to scale down the image size into half. By converting to .npy format can also increase the loading speed:

python convert.py -i data/tokyo.jpg -o data/tokyo.npy --scale 0.5

Then run the experiment

python train.py -i data/tokyo.npy --enc_method hashgrid --finest_resolution 32768 --visualize

hash-grid-encoding's Issues

Performance between the implementation and the original version

Thanks for your work! I know it is difficult to transfer hash encoding to pure python implementation. I am wondering the performance (i.e. PSNR, rendering quality) between this implementation and the original cuda/c++ version. Have you tested it on NeRF synthetic dataset or other data? Thanks!