search

threedle / text2mesh

3D mesh stylization driven by a text input in PyTorch
https://threedle.github.io/text2mesh/
MIT License
925 stars 129 forks source link

[FIX] Improve model saving #34

Closed ruiqixu37 closed 5 months ago

ruiqixu37 commented 6 months ago

Hi,

Thanks for the amazing work! I am working on an extended object based on text2mesh and found that the current save_model function in neural_style_field.py doesn't quite work. This is because the intermediate variables for the ProgressiveEncoding and FourierFeatureTransform layers are not saved. These variables have critical impacts on model performances and will lead to inconsistent performance if not saved.

The solution I find is pretty simple: just register _t in ProgressiveEncoding, and _B in FourierFeatureTransform as nn.Parameters. They will then be saved by PyTorch's model.state_dict() correctly. It works for me when I try to load a saved model and conduct inference now.

Think it is worth sharing it here if anyone has met similar issues. If you think the above makes sense, I can open a pull request to fix this.

Love to hear your thoughts. Thanks!

factoryofthesun commented 5 months ago

Hi sorry for the delayed response and thanks for the proposed fix! We had a similar fix in our project code, but it seems that didn't end up making it over to this public repo. Please go ahead and open the pull request we'd love to be able to make sure that the model results are reproducible. Thanks.

ruiqixu37 commented 5 months ago

Thanks for the reply. Will do!