[UnitaryHack] Make ansatz class a training hyperparameter

[neiljdo]

PR adds the ansatz class and necessary arguments for instantiation (e.g. ob_map and ansatz-specific kwargs) as hyperparameters to the trainer class. The functions related to creating and loading checkpoints have also been updated to accommodate the new ansatz class hyperparameter.

Addresses #86.

[ianyfan]

Hi Neil, thanks for your submission, it looks quite good already. I've enabled the tests to run on this PR, so would you be able to have a look at where they're failing and see if you could fix them up?

[neiljdo]

Hi @ianyfan, thank you for the review. I just pushed a new commit to address the flake8 violations. I'll wait for the workflow run results to see if I missed some.

[neiljdo]

@ianyfan just saw the failing jobs related to the notebook tests and type checking - will push a fix to those soon.

[dimkart]

@ianyfan just saw the failing jobs related to the notebook tests and type checking - will push a fix to those soon.

@neiljdo Hi, it would be great if you could that by tomorrow, 13/6, which is the last day of the hackathon.

[nathanshammah]

@dimkart @neiljdo PRs can be reviewed also shortly after June 13th and still be awarded a bounty in uHACK if accepted.

[dimkart]

@nathanshammah

PRs can be reviewed also shortly after June 13th and still be awarded a bounty in uHACK if accepted.

Thanks, noted.

[nikhilkhatri]

Hi @neiljdo ,
Thanks for this PR!
it looks good in general, though I have some comments:

1. The ansatz is applied to the diagrams at each epoch. This can slow down training a lot, especially for large datasets / diagrams. Instead, it would be better to save the train and val circuits as properties of the Trainer object during init and checkpoint-load.
2. It would be good to add tests / manually test that all the ansatze are serialisable through pickle.
3. It would be much better if the model did not take the test set as input during training, for a clear separation. This would be particularly natural once the UNK feature is completed. (UnitaryHACK: Replace unknown words in diagrams with UNK token #84)
4. Currently, the ansatz is a property of the Trainer, though it makes conceptual sense for this to be a property of the Model. In this case, it would be possible to make the forward and get_diagram_output functions take a diagram as input, and apply the ansatz internally.

Since this is already a complex task and there are some design decisions involved, it would be great if you could address points 1 and 2 so the issue can be assigned to you for the hack.

The remaining points can be discussed later.

[nikhilkhatri]

Would also be great if you could run the clean_notebooks.py script, which removes some metadata from the examples and tutorials.

[neiljdo]

@nikhilkhatri,
Thank you for the review. Here are my thoughts on your comments above.

1. We should save the circuits in order to do the diagram-to-circuit conversion only once. I'll implement a fix for this one right away.
2. For the ansatz serialization, I only save the components (i.e. the ansatz class, ob_map, and other ansatz-specific kwargs) required to initialize the ansatz during trainer checkpoint creation and re-initialize the ansatz during checkpoint load. I thought the ansatz deterministic and could be re-instantiated, so I went this route. I would like to do tests as part of the testing suite (not manually) - this implies adding new tests that use different ansatzes for the trainers.

For (3) and (4) above, I realized that the model is coupled with the circuits while working on the issue. I saw this more while updating the demo notebooks and saw that the Model.from_diagram call also needs to get the test circuits. Let's discuss these two later once I've resolved the first two.

[nikhilkhatri]

@neiljdo you are correct in that the ansatze are deterministic. My point was to ensure that there are no Exceptions raised when pickling an ansatz class (for example, if they use a lambda function).

[nikhilkhatri]

w.r.t the lambda pickling, your code may be alright since it pickles the class, not the instance. It would still be good to make sure.

[ACE07-Sev]

May I drop a suggestion here? Based on what I understand from the code :

1. If Model is from checkpoint, so you just skip the _init_model_from_datasets method , but still have to recreate the circuits at training_step and validation_step (assuming number 1 is fixed).

2. If Model is not from checkpoint, you make the circuits at both the _init_model_from_datasets method and the training_step and validation_step.

Since the dataset method is not run every time, I think it could be better if you make the circuits once, pass the list to _init_model_from_datasets method, and also use the list in the training and validation step (based on the index of the diagram of course).

My suggestion is after fixing 1, to have the _init_model_from_datasets also have the reference from 1's fix. @nikhilkhatri dear, am I correct? Please correct me if I am wrong.

[nikhilkhatri]

Hi @ACE07-Sev ,
Thanks for your comments. Indeed, the circuits should only be created once, and then used wherever the ansatz is currently applied. The circuits need to be created wherever the dataset and ansatz parameters are first available. If the dataset is taken as input first during fit, then this is where the ansatz can be applied, and circuits from this can be reused throughout the setup, training and validation.

[neiljdo]

Hi @nikhilkhatri,

I've addressed points (1) and (2) from your original comment in my latest update. Specifically,

1. Circuits are now generated once at the start of the Trainer.fit() execution. Circuits are also included in the checkpoint.
2. I've used dill as a drop-in replacement for pickle for checkpoint generation as it handles lambda functions well. I've added tests for all available ansatzes to make sure that it still generates the correct output after serialization-deserialization. With this, we can simplify the ansatz checkpoint creation i.e. no need to re-instantiate.

Let me know if there are still issues, thank you very much!

EDIT: I ran the clean_notebooks.py script but it did not update the notebooks.

[nikhilkhatri]

Hi @neiljdo
Thanks a lot for these changes!

I have a couple of comments:

1. Saving all the circuits will create a very large checkpoint file, depending on the dataset size, and circuit complexity. I dont think this is required. If a trainer is loaded from a checkpoint, the user will likely call fit, at which point they can provide the necessary circuits.
2. The tests seem to be failing because dill hasnt been added to requirements.
3. (Minor) The tests for pickling the ansatze are very verbose. Using something like this can help here.

It would be great if you could address issues 1 and 2 for the hack.

[neiljdo]

@nikhilkhatri
Noted, thanks. I'll remove the circuits from the checkpoint.

[ACE07-Sev]

Neil dear :

lambeq/training/checkpoint.py:27: error: Cannot find implementation or library
stub for module named "dill" [import]

The test with pytest passed, so I think you just need to fix this import error.

[dimkart]

@neiljdo Hey Neil please leave a comment in the Issue #86, so you can be assigned to it (otherwise your handle cannot be found).

[dimkart]

Thanks, good work! The issue will be assigned to you, however the PR will be merged a little later.

[dimkart]

Moved to lambeq's private repository for further development. This feature will be included in one of the upcoming releases.