Add UKAEA-TGLFNN network #11
Conversation
|
FYI, we have just renamed the library to fusion_surrogates, and this will cause import issues with your changes.
Apologies for the trouble. |
theo-brown
force-pushed
the
ukaea-tglfnn
branch
from
618cb81 to
b00e087
Compare
|
This is complete pending the release of the model weights and/or data for testing. I could also plausibly add more unit tests. |
|
Thanks for the review @hamelphi. It was a bit of a half-baked PR because things were changing internally and we weren't sure about the methods we'd use for exporting/sharing these models. Things should be a bit more consistent now, following your comments! |
There was a problem hiding this comment.
Please move the tests associated to these transforms outside from qlknn_model_test.py into transforms_test.py
| efe_onnx_path: str, | ||
| efi_onnx_path: str, | ||
| pfi_onnx_path: str, | ||
| ) -> "TGLFNNModel": |
There was a problem hiding this comment.
Remove this type annotation.
There was a problem hiding this comment.
Now that the loaders are in different files, we should split these tests as well.
| model = pytorch_model.PytorchTGLFNNModel( | ||
| config_path="models/1.0.1/config.yaml", | ||
| stats_path="models/1.0.1/stats.json", | ||
| efe_gb_checkpoint_path="models/1.0.1/regressor_efe_gb.pt", |
There was a problem hiding this comment.
We should add the model data and test data in this PR as well. My understanding is that you are waiting for the approval on your side before pushing the data. Is that correct?
|
FYI, I refactored the file structure to follow what you suggested in this PR. I added the |
Adding a refactored version of the UKAEA-TGLFNN from google-deepmind/torax#477. I've tried to match the format of the model file to the qlknn model - feedback welcome!
Tasks: