Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Functional laplace #55

Closed
wants to merge 94 commits into from
Closed

Functional laplace #55

wants to merge 94 commits into from

Conversation

metodmove
Copy link
Contributor

No description provided.

metodj and others added 30 commits September 15, 2021 22:07
@metodj
init commit
e3f0481
@metodmove
FunctionalLaplace blueprint
bcb4bfb
@metodmove
added a test that will serve as a sanity check for correctness
b2a2ff4
@metodmove
test small refactor
a793a04
@metodmove
init GP matrices and SoD dataloader
1b50f1d
@metodmove
initial full naive GP implementation
84bd1a7
@metodmove
fixed bugs in full naive GP inference
8050bff
@metodmove
test gp_equivalence (almost) passing
2ac2b81
@metodmove
add GP to regression_example
cdbadae
@metodmove
float32 vs float64 magic
3e4ffd2
@metodmove
Cholesky solves float32 vs float64 conundrum
11943c1
@metodmove
minor
e43eb0e
@metodmove
merge remove-abc into functional-laplace
8b51f62
@metodmove
fix multivariate regression GP bug
aab27a6
@metodmove
GP inference parameters
583cef6
@metodmove
GP classification with diagonal_L=True
cf00987
@metodmove
start independent_gp_kernels=True
0e28c4a
@metodmove
block-diagonal kernel analysis
d28cc6b
@metodmove
warning for multivariate regression
4b568cb
@metodmove
refactor kernel method
4f58447
@metodmove
prior factor for SoD and refactor __call__
bff30f4
@metodmove
log_marginal_likelihood GP start
d61b4fa
@metodmove
scatter for logp(f)
3a5a1be
@metodmove
fix gp marginal likelihood bugs
dec9514
@metodmove
GP predictive_samples
4dc7016
@metodmove
start with tests
7ea6328
@metodmove
minor
c91048e
@metodmove
more tests and refactor
7376204
@metodmove
docstrings
ab5058a
@metodmove
add last layer functional laplace
0ef5ae3
Metod Jazbec and others added 20 commits December 30, 2022 11:50
further optimize delta experiment
6c7397e
run for larger delta
92bfacf
last-layer debug
4b08edb
inference speed-up
87abeeb
minor
7747168
einsum memory investigation
09d410a
CV working
47eebfe
rebuild on Sigma_inv
648b3ed
clean
eef9a15
validate no_grad
8967f9d
@metodj
Functional laplace gp calibration (#1)
edf6285 
* gp calibration notebook

* minor

* inducing points CIFAR experiment

* transfer model from bnn-preds repo

* inducing points FMNIST CNN

* gp calibration example

* gp continue

* subset_of_weights=all experiment

* fixed prior precision

* fixed prior precision

* ensure that input is differentiable

* further optimize delta experiment

* run for larger delta

* last-layer debug

* inference speed-up

* minor

* einsum memory investigation

* CV working

* rebuild on Sigma_inv

* clean

* validate no_grad

Co-authored-by: Metod Jazbec <metodjazbec@Metods-MacBook-Pro.local>
Co-authored-by: Metod Jazbec <metodjazbec@wcw-staff-145-109-82-97.wireless.uva.nl>
Co-authored-by: Metod Jazbec <metodjazbec@uvavpn-byodm-145-18-162-61.vpn.uva.nl>
Co-authored-by: Metod Jazbec <metodjazbec@uvavpn-byodm-145-18-160-65.vpn.uva.nl>
Co-authored-by: Metod Jazbec <metodjazbec@uvavpn-byodm-145-18-162-52.vpn.uva.nl>
Co-authored-by: Metod Jazbec <metodjazbec@wcw-staff-145-109-84-166.wireless.uva.nl>
Co-authored-by: Metod Jazbec <metodjazbec@wcw-staff-145-109-89-247.wireless.uva.nl>
fix tests
edd83b1
Merge branch 'functional-laplace-gp-calibration' into functional-laplace
88b44d7
clean calibration-gp example
44fb5c2
minor
c9c068d
minor2
39000ae
markdown example start
5cb1354
markdown wrapup
d8ec55b
minor
03f0989
minor
005c3c7
@metodj
Copy link
Contributor

metodj commented Jan 12, 2023

Main changes since last time:

  • addressed underfitting of FunctionalLaplace in regression_example.py, fixed the issue with the marginal likelihood: Functional laplace #55 (comment)
  • added calibration_gp_example.py where FunctionalLaplace is used on a pre-trained FMNIST classifier (same experiment as in linearised Laplace paper)

@metodj
Copy link
Contributor

metodj commented Jan 12, 2023

I also benchmarked the code with BNN-predictions-repo in terms of speed and the code here is 2-3x slower. After investigating, I believe the main reason is that code here allows one to only use smaller batch sizes before running into memory issues when doing GP inference (b=128 here compared to e.g. b=512 in the BNN-predictions-repo).

When it comes to memory, the bottleneck seem to be torch.einsum functions used in the kernel methods in FunctionalLaplace. To get around this, I also tested implementation without torch.einsum. However, then the memory issues started to pop up in the backpack backend.

Further benchmarking and faster implementation can be left for future PRs imo.

metodj and others added 4 commits February 21, 2023 15:22
@metodj
Functional laplace memory investigation (#2)
e20ac6c 
* memory investigation start

* add for loop over output dimensions

* improve memory footprint

* minor

---------

Co-authored-by: Metod Jazbec <mjazbec@ivi-cn011.ivi.local>
@metodj
Functional laplace memory investigation (#3)
35ec107 
* memory investigation start

* add for loop over output dimensions

* improve memory footprint

* minor

* fix tests

---------

Co-authored-by: Metod Jazbec <mjazbec@ivi-cn011.ivi.local>
@metodj
increase batch size in the example
77a32f4
@metodj
minor
5de3980
@metodj
Copy link
Contributor

metodj commented Feb 21, 2023 •
edited
Loading

Improved memory performance (see changes in _kernel methods in FunctionalLaplace and in jacobians method in BackPackInterface). Can use larger batch sizes (e.g. b=256 on Nvidia 3090 GPU) now which results in around 2x speedup, so more or less on par with implementation in BNN-predictions-repo

@Ludvins Ludvins mentioned this pull request Jun 4, 2024
Merged
@wiseodd
Copy link
Collaborator

wiseodd commented Jun 23, 2024

Replaced by #192

@wiseodd wiseodd closed this Jun 23, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@metodj metodj metodj left review comments

@aleximmer aleximmer aleximmer left review comments

Assignees
No one assigned
Labels
enhancement New feature or request
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
4 participants
@metodmove @aleximmer @metodj @wiseodd