✨ LeNet BatchEnsemble and Deep Ensemble + minor bugfixes

[tonyzamyatin]

This PR introduces the following improvements and fixes:

• Bug Fixes: Address minor issues in some of the existing LeNet configs.
• BatchEnsemble Wrapper: Implements a wrapper that replicates the input batch num_estimator times, eliminating the need to hardcode this for each model.
• Documentation Updates: Enhances the documentation to reflect the usage of batch ensemble layers and the classification routine, aligning with available wrappers and transforms.
• LeNet BatchEnsemble Implementation: Adds BatchEnsemble support for LeNet, along with new experiment configurations for both Deep Ensemble and BatchEnsemble models (tested manually).

[o-laurent]

Hello again @tonyzamyatin,

We have long wanted to do a BatchEnsemble wrapper, so thank you for this. Before I get to looking at your PR in detail, don't hesitate to enable the pre-commits as detailed in the contributing page. By the way, you could directly enable them in your dockerfile in #134?

[tonyzamyatin]

Hi @o-laurent,

currently I have a fork with where I play around and if I implement something new or fiy a bug I make a new branch and create a PR.

I have the pre-commit checks enabled on my fork and manually run the hooks before I make commits.

Are there any problems, or ways to automate this, just because you're asking?

And yeah I can add that to the Dockerfile👌

[o-laurent]

No worries, I'm saying this to avoid failing GitHub workflows due to linting issues as in test #1601.

If the docker image is made for contributions, then I think that it makes sense to enforce that the pre-commit hooks are installed. When installed, they should run automatically at commit time - and not manually - unless you use the --no-verify flag.

[tonyzamyatin]

Yeah you're right. Definitely, makes sense to enable pre-commit hooks in the Dockerfile.

Hmmm when I commit something onto the fork the pre-commit hook doesn't automatically run, only the pre-commit checks.🧐

[o-laurent]

Hmmm when I commit something onto the fork the pre-commit hook doesn't automatically run, only the pre-commit checks.🧐

What do you mean by pre-commit check (vs. pre-commit hooks)?

Normally, pre-commit should make committing impossible when any of the hooks fail, in this case "ruff-check".

[tonyzamyatin]

Sorry for the unclarity.

The commits get rejected if the pre-commit hook doesn't successfully run, but the code isn't automatically formatted on commit. I have to run the pre-commit hook manually and then commit again.

[o-laurent]

On my side, I just have to add the modified files and commit again. Anyway, this doesn't explain how you managed to push code that made our test fail 😄 but it's no big deal at all.

[tonyzamyatin]

Exactly, I do the same.

Don't know how this happens lol 😂

[tonyzamyatin]

Altough, usually it's some kind of test for a transform which is not implemented that makes the pipeline fail.

This time I maybe forgot to lint🌚

I'll have to double check whether the pre-commit hooks is active in the Docker container by default.

[o-laurent]

Damn, I didn't install the pre-commits on my clone of your fork and now I'm pushing failing code too 🤣

[tonyzamyatin]

HAHAHHA🤣

[codecov]

Codecov Report

Attention: Patch coverage is 98.36066% with 1 line in your changes missing coverage. Please review.

Project coverage is 99.15%. Comparing base (7d0d345) to head (c737d07).
Report is 62 commits behind head on dev.

✅ All tests successful. No failed tests found.

Files with missing lines	Patch %	Lines
...orch_uncertainty/models/wrappers/batch_ensemble.py	97.22%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##              dev     #137      +/-   ##
==========================================
- Coverage   99.15%   99.15%   -0.01%     
==========================================
  Files         143      144       +1     
  Lines        7136     7192      +56     
  Branches      910      923      +13     
==========================================
+ Hits         7076     7131      +55     
- Misses         27       28       +1     
  Partials       33       33              

Flag	Coverage Δ
cpu	99.15% <98.36%> (-0.01%)	⬇️
pytest	99.15% <98.36%> (-0.01%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[o-laurent]

Thank you for your PR! Here are a few quick comments on some details.

[alafage]

Hi @tonyzamyatin,

Nice PR, thank you! I don't know if you saw my answer on Discord, but there are some points that we need to address before merging these changes to dev.

First, we would like to keep the possibility of not repeating the batch during training to fit the original paper better. In this case, we do not need to use format_batch_fn=RepeatTarget(num_repeats=num_estimators) in the ClassificationRoutine().

Secondly, @o-laurent and I agree that the current BatchEnsemble wrapper can be misleading. It does not make a BatchEnsemble but overrides the model's forward to repeat the inputs first. It would be clearer if we had a wrapper taking any nn.Module and dynamically replacing the layers with their BatchEnsemble layer counterparts.

I will make these changes directly on your fork if that's okay.

[tonyzamyatin]

Hi @alafage,

That makes total sense!

First, we would like to keep the possibility of not repeating the batch during training to fit the original paper better. In this case, we do not need to use format_batch_fn=RepeatTarget(num_repeats=num_estimators) in the ClassificationRoutine().

For this, I think we could simply pass a boolean flag via the __init__() constructor to indicate whether to replicate the targets for training or not.

Secondly, @o-laurent and I agree that the current BatchEnsemble wrapper can be misleading. It does not make a BatchEnsemble but overrides the model's forward to repeat the inputs first. It would be clearer if we had a wrapper taking any nn.Module and dynamically replacing the layers with their BatchEnsemble layer counterparts.

This also makes a lot of sense, but note that this introduces additional complexity, could lead to silent failures if there is no such implementation, and cause unintended behaviour (e.g. using a pretrained model with the wrapper). Maybe its more explicit to implement BatchEnsemble constructors similar to lenet_batchensemble() in lenet.py?

In any case, I belive it makes sense to add a Note in the docstring that the architecture must be specified to use BatchEnsemble layers.

What do you think?

[tonyzamyatin]

All feedback implemented apart from this detail., but I can also make this change if you want :)

[alafage]

I added a boolean repeat_training_inputs in torch_uncertainty.wrappers.BatchEnsemble.__init__() to control whether to repeat the inputs during training. The inputs are always repeated for evaluation.

Concerning the second point, although it adds more complexity, I have added another boolean parameter to the class: convert_layers. If True, all nn.Linear and nn.Conv2d layers will be replaced by BatchLinear and BatchConv2d layers. To ensure the wrapper is correctly used, we check if there is at least one BatchLinear or BatchConv2d layer in the model. If not, an error will be raised. I took the liberty to update bachensemble_lenet accordingly.

Although I've written tests to validate my modifications, feel free to check that they do not break anything for your experiments.

Cheers!

[o-laurent]

Seems perfect! Thanks @tonyzamyatin and @alafage!