NN Extraction For Conversion and use in Other Formats

[Davicam2]

Outstanding project scientists!

This project achieves wonderful results Orders of magnitudes faster than anything else out there.

I would like to take the trained model, and apply its use in other engines/projects, but am not sure how to extract the NN. the .msgpack as far as i can tell is built for the test bed, and i cant seem to find the trained NN being stored anywhere.

Please let me know if this is a future development or if i can manipulate the test with python to extract a trained model!

[collinabidi]

^ I'm interested in this as well.

I assume that we will have to use the tiny-cuda-nn package like this, but we don't have a .json.

[Tom94]

Hi there, neither instant-ngp now tiny-cuda-nn directly support exporting models in portable formats.

That said, you can replicate the computation graph from NerfNetwork in your favorite machine learning framework and obtain the raw trained parameters either by

• reading the MessagePack file (which is essentially a binary json-equivalent that has implementations in all kinds of languages) and extracting the snapshot.params_binary field (can be fp16 or fp32, depending on config / machine), or
• by calling m_trainer->params() within Testbed to get pointer to GPU memory holding the trained parameters in fp32 format.

NerfNetwork::set_params and callees showcase how the parameters are laid out. Note that the underlying MLP implementations CutlassMLP and FullyFusedMLP apply different amounts of padding to the last layer (8 vs. 16).

Cheers!

[HenryBarrett12]

Hi there, excellent work!
I am working on this for a University project and have a similar question as the original commenter.

I was wondering if you could elaborate on the above process using PyTorch bindings?

Currently, I am able to initalise a tiny-cuda-nn model using the snapshots encoding and network config, but am yet to load the params_binary into the model. What steps do I need to take to manually load these parameters?

My guess is that some manipulation is required to convert the raw parameters into a model_dict format that can be loaded with pytorch load_state_dict but I am unsure how to do so.

My current process is as follows:

1. Parse msgpack into json format
2. Create model tcnn.NetworkWithInputEncoding(n_input, n_output, config["encoding"], config["network"])
3. ? - load network parameters
4. Query model as usual

Any advice or code snippets would be greatly appreciated!

[rjeye]

@HenryBarrett12 Did you have any success with your proposed method ? Any learnings ?

[kwea123]

To extract the weights from .msgpack, the following code works:

import msgpack

with open('transforms_train_base.msgpack', 'rb') as f:
    data = msgpack.loads(f.read())

then data is a dictionary containing model structure and weights (in binary). For example, to load the density grid, you can do

density_grid = np.frombuffer(data['snapshot']['density_grid_binary'], dtype=np.float16)

and same for network weights. Next, to assign these values to a tinycudann I think is easier

[robbie1383]

Hi!

I am trying to follow these steps. I trained my network and I got the .msgpack file out of it. I can read it with the code snippets you provided but that is where I get stuck.

I see I need to use tcnn.NetworkWithInputEncoding(n_input, n_output, config["encoding"], config["network"]), but I don't know what my n_input and n_output are. I also am not sure how to put in the weights into the network afterwards or where to find all this information in the data dictionary.

Do you know of any tutorials or code snippets or anything that could help?

[fedeceola]

Hi @Tom94, is there a function to extract only the hash encoding of a given point in the grid (the y vector in fig 3-(4) of the paper)?

@HenryBarrett12, @rjeye, did you manage to solve the problem with the pytorch tiny-cuda-nn bindings?