Add a range of features for augmentation and training, and fix a few bugs

[chris-rapson-hectre]

Motivation

Improve augmentation and training for our model, and hopefully others' models as well.

Modification

• generalise the heatmap loss calculated in ae_head to have the correct dimensions
• new augmentations: bidirectional flip and rotation by Nx90º
• visualisation to debug augmentations
• add a hook to initially freeze and then unfreeze the backbone, for safer fine tuning
• add a hook for early stopping
• bug fix for singular Hessian matrices (e.g. a constant heatmap)

BC-breaking (Optional)

No

Use cases (Optional)

1. associative embedding training failed for me because of a shape mismatch with the keypoint_weights
2. more augmentations are always helpful
3. verifying that augmentations are working as expected helps to debug algorithms and tune hyperparameters
4. pretrained backbones can be vulnerable to large losses during the first few iterations when fine tuning, so it's good to freeze them initially and then allow them to adapt later
5. stop the training when the model starts to overfit to avoid wasting GPU resouce
6. using pinv instead of inv means the typical case is unchanged and avoids crashing on a pathological edge case.

Checklist

Before PR:

• [y ] I have read and followed the workflow indicated in the CONTRIBUTING.md to create this PR.
• [ y] Pre-commit or linting tools indicated in CONTRIBUTING.md are used to fix the potential lint issues.
• [ y] Bug fixes are covered by unit tests, the case that causes the bug should be added in the unit tests.
• [y ] New functionalities are covered by complete unit tests. If not, please add more unit tests to ensure correctness.
• [n ] The documentation has been modified accordingly, including docstring or example tutorials.

The only documentation in this PR is as docstrings. Please let me know if you want more.

After PR:

• CLA has been signed and all committers have signed the CLA in this PR.

[CLAassistant]

All committers have signed the CLA.

[chanul-hectre]

Indent the else block.

[chanul-hectre]

These lines were added in other repo:

Suggested change

	model = runner.model
	assert hasattr(runner, 'epoch'), 'Runner must have an epoch attribute'
	assert hasattr(runner.model, 'backbone'), 'Model must have a backbone attribute'
	model = runner.model