Skip to content

Jittable gnn

Xiangyang Ju requested to merge xju-mr-jittable into dev
  • Adding two variations of the InteractionGNN2: RecurrentInteractionGNN2 & ChainedInteractionGNN2. Such division is necessary to make the GNN jittable.
  • Adding a script that converts a trained model to TorchScript and ONNX. This script works with Metric Learning and Filtering models too. Try it with python scripts/save_full_model.py gnn_train.yaml or python scripts/save_full_model.py gnn_train.yaml -c your-model-ckpt.
  • Adding a notebook that tests if the conversion work.
Edited by Xiangyang Ju

Merge request reports