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.
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')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.
Run this example:
python train.py -i data/albert.jpg --enc_method hashgrid --visualize
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.5Then run the experiment
python train.py -i data/tokyo.npy --enc_method hashgrid --finest_resolution 32768 --visualize