Loading weights into tiny-cuda-nn of C++/CUDA from pre-trained PyTorch model

[SJoJoK]

Thanks for the amazing work! I'm using tiny-cuda-nn in both python and CUDA way. The issue I met is that I cant find a way to save the model parameters trained by the pytorch extension and load them into CUDA. Previous discussion https://github.com/NVlabs/tiny-cuda-nn/issues/6 seems doesnt related to the pytorch extension. Thanks in advance.

[SJoJoK]

UPDATE: I found that state_dict['params'] of tcnn.NetworkWithInputEncoding is a long tensor, which maybe the "densely laid out network parameters" described in #6. I'll try to directly load it into CUDA, and i will close the issue if I success.

[robbie1383]

Hi!

I am currently stuck on the same part. Did you have any luck with importing the weights?

[SJoJoK]

Hi!

I am currently stuck on the same part. Did you have any luck with importing the weights?

Hello @robbie1383 , I indeed successed somhow months ago. state_dict['params'] of tcnn.NetworkWithInputEncoding is a long tensor, you might check https://github.com/pytorch/pytorch/issues/36568 and https://discuss.pytorch.org/t/load-tensors-saved-in-python-from-c-and-vice-versa/39435/8. to load it into C++/CUDA, and then use Trainer::set_params to set the params. However, I only checked whether the params can be loaded in C++/CUDA, but not checked if the outputs are correct. Hopes these can help you.

[zyc-bit]

Hi, is there bias in your state_dict? Or your network purely uses weight matrices + activations?

[SJoJoK]

Hi, is there bias in your state_dict? Or your network purely uses weight matrices + activations?

Sorry I can't help you because the issue was created long time ago and I almost forget all the relative information.