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.

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

Allow for matrix-matrix and matrix-vector products with KFACLinearOperator and KFACInverseLinearOperator without converting to numpy #91

Merged
merged 16 commits into from
Mar 20, 2024

Conversation

runame
Copy link
Collaborator

@runame runame commented Mar 15, 2024

Addresses #71, but only for the KFACLinearOperator and KFACInverseLinearOperator.

KFACLinearOperator/KFACInverseLinearOperator is arguably the only linear operator here that is likely to be used for preconditioned-gradient methods for large scale neural networks. Therefore, adding torch_matmat and torch_matvec methods seem like a simple solution to avoid unnecessary device transfers, which are a bottleneck for this use case. However, this doesn't address the issue in general.

@runame runame added the enhancement New feature or request label Mar 15, 2024
@runame runame requested a review from f-dangel March 15, 2024 14:46
@runame runame self-assigned this Mar 15, 2024
Copy link
Contributor

@sourcery-ai sourcery-ai bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hey @runame - I've reviewed your changes and they look great!

General suggestions:

  • Consider refactoring input validation and device compatibility checks into separate methods to improve code readability and maintainability.
  • Explore more efficient approaches for tensor concatenation to optimize performance.
  • Enhance error messages with more specific details to aid in debugging and improve the developer experience.
Here's what I looked at during the review
  • 🟡 General issues: 4 issues found
  • 🟢 Security: all looks good
  • 🟢 Testing: all looks good
  • 🟢 Complexity: all looks good
  • 🟢 Docstrings: all looks good

Thanks for using Sourcery. We offer it for free for open source projects and would be very grateful if you could help us grow. If you like it, would you consider sharing Sourcery on your favourite social media? ✨

Share Sourcery

Help me be more useful! Please click 👍 or 👎 on each comment to tell me if it was helpful.

curvlinops/kfac.py Outdated Show resolved Hide resolved
curvlinops/kfac.py Outdated Show resolved Hide resolved
Comment on lines 346 to 347
if return_tensor:
M_torch = cat([rearrange(M, "k ... -> (...) k") for M in M_torch], dim=0)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

suggestion (performance): Consider the efficiency of tensor concatenation.

Concatenating tensors in a loop can be inefficient, especially for large numbers of tensors. It might be beneficial to explore alternative approaches that could reduce the computational overhead, such as preallocating a tensor of the correct size and filling it.

curvlinops/kfac.py Outdated Show resolved Hide resolved
curvlinops/kfac.py Outdated Show resolved Hide resolved
Co-authored-by: sourcery-ai[bot] <58596630+sourcery-ai[bot]@users.noreply.github.com>
@coveralls
Copy link

coveralls commented Mar 15, 2024

Pull Request Test Coverage Report for Build 8360156180

Details

  • 52 of 58 (89.66%) changed or added relevant lines in 2 files are covered.
  • No unchanged relevant lines lost coverage.
  • Overall coverage increased (+0.2%) to 89.367%

Changes Missing Coverage Covered Lines Changed/Added Lines %
curvlinops/inverse.py 16 17 94.12%
curvlinops/kfac.py 36 41 87.8%
Totals Coverage Status
Change from base Build 8250604944: 0.2%
Covered Lines: 1017
Relevant Lines: 1138

💛 - Coveralls

@runame runame changed the title Allow for matrix-matrix and matrix-vector products with KFACLinearOperator without converting to numpy Allow for matrix-matrix and matrix-vector products with KFACLinearOperator and KFACInverseLinearOperator without converting to numpy Mar 15, 2024
Copy link
Owner

@f-dangel f-dangel left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I made some refactoring proposals.

curvlinops/inverse.py Outdated Show resolved Hide resolved
curvlinops/inverse.py Outdated Show resolved Hide resolved
curvlinops/inverse.py Outdated Show resolved Hide resolved
curvlinops/inverse.py Outdated Show resolved Hide resolved
curvlinops/inverse.py Outdated Show resolved Hide resolved
curvlinops/kfac.py Outdated Show resolved Hide resolved
curvlinops/kfac.py Outdated Show resolved Hide resolved
curvlinops/kfac.py Outdated Show resolved Hide resolved
curvlinops/kfac.py Outdated Show resolved Hide resolved
test/test_kfac.py Outdated Show resolved Hide resolved
@runame runame requested a review from f-dangel March 18, 2024 17:35
Copy link
Owner

@f-dangel f-dangel left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Minor things. After this round I believe we are ready to merge, maybe up to some more minor nits.

curvlinops/inverse.py Outdated Show resolved Hide resolved
curvlinops/inverse.py Outdated Show resolved Hide resolved
curvlinops/kfac.py Outdated Show resolved Hide resolved
curvlinops/kfac.py Outdated Show resolved Hide resolved
curvlinops/kfac.py Outdated Show resolved Hide resolved
test/test_inverse.py Outdated Show resolved Hide resolved
test/test_kfac.py Outdated Show resolved Hide resolved
test/test_kfac.py Show resolved Hide resolved
test/test_kfac.py Outdated Show resolved Hide resolved
test/test_kfac.py Outdated Show resolved Hide resolved
Copy link
Owner

@f-dangel f-dangel left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Found some more minor things, but looks great! One more thing to discuss is whether we want to switch to using cases in the tests of KFAC.

curvlinops/kfac.py Outdated Show resolved Hide resolved
curvlinops/inverse.py Outdated Show resolved Hide resolved
curvlinops/kfac.py Outdated Show resolved Hide resolved
test/test_inverse.py Show resolved Hide resolved
test/test_inverse.py Show resolved Hide resolved
test/test_inverse.py Outdated Show resolved Hide resolved
test/test_inverse.py Outdated Show resolved Hide resolved
test/test_kfac.py Outdated Show resolved Hide resolved
@runame runame requested a review from f-dangel March 20, 2024 13:51
Copy link
Owner

@f-dangel f-dangel left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks ready to me!

@runame runame merged commit 67e043f into main Mar 20, 2024
7 checks passed
@runame runame deleted the kfac-torch-only branch March 20, 2024 19:23
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
enhancement New feature or request
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants