search

RaduAlexandru / permutohedral_encoding

Blazingly fast encoding for neural networks based on permutohedral lattices
MIT License
97 stars 10 forks source link

Remove internal C++ members from state dictionary to allow saving #2

Closed LvisRoot closed 1 year ago

LvisRoot commented 1 year ago

Looks like internal C++ bind members are not serializable by pickle so it fails when trying to save a model using the PermutoEncoder with the following error:

TypeError: cannot pickle 'permutohedral_encoding_bindings._86_C.EncodingFixedParams' object

This PR removes them from the dictionary before pickling and rebuilds them when loaded.

Took the idea from tinycudann way of saving their encodings (src)

hope this make sense and I didn't miss another way to save the models.

RaduAlexandru commented 1 year ago

Thanks a lot for the PR!

Indeed, saving the model using torch.save(encoding, PATH) seems to fail with the current setup because of the C++ members. I think I didn't encounter this issue in my experiments since I was using torch.save(encoding.state_dict(), PATH) which just saves the learnable parameters. The official docs have more info about their differences. Either way, both methods of saving should be supported.

Thanks again for catching this and implementing the fix!