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

Refactor dual cone projections #237

Merged

ValerianRey merged 36 commits into main from factorize_dual_cone_projections

Feb 4, 2025

+164 −94

Contributor

PierreQuinton commented Jan 21, 2025

Both DualProj and UPGrad rely on projections onto the dual cone, this PR factorizes the computation of such projections into one utility in the aggregation package.

PierreQuinton added the package: aggregation label

PierreQuinton requested a review from ValerianRey

January 21, 2025 15:31

PierreQuinton self-assigned this

codecov bot commented Jan 21, 2025 •

edited

Loading

Codecov Report

All modified and coverable lines are covered by tests ✅

Files with missing lines	Coverage Δ
src/torchjd/aggregation/_dual_cone_utils.py	`100.00% <100.00%> (ø)`
src/torchjd/aggregation/_gramian_utils.py	`90.90% <100.00%> (+3.40%)`	⬆️
src/torchjd/aggregation/dualproj.py	`100.00% <100.00%> (ø)`
src/torchjd/aggregation/upgrad.py	`100.00% <100.00%> (+1.96%)`	⬆️

PierreQuinton added the refactor label

Base automatically changed from explicit_dtypes to main

January 21, 2025 16:24

PierreQuinton force-pushed the factorize_dual_cone_projections branch from d9f60e3 to 5c15741 Compare

January 21, 2025 17:15

Contributor

ValerianRey commented Jan 27, 2025

I really like the idea! I'll make a proper review and probably merge later this week.

ValerianRey reviewed

View reviewed changes

src/torchjd/aggregation/_dual_cone_utils.py Outdated Show resolved Hide resolved

ValerianRey reviewed

View reviewed changes

src/torchjd/aggregation/_dual_cone_utils.py Outdated Show resolved Hide resolved

ValerianRey reviewed

View reviewed changes

tests/unit/aggregation/test_dual_cone_utils.py Outdated Show resolved Hide resolved

ValerianRey reviewed

View reviewed changes

tests/unit/aggregation/test_dual_cone_utils.py Outdated Show resolved Hide resolved

ValerianRey reviewed

View reviewed changes

tests/unit/aggregation/test_dual_cone_utils.py Outdated Show resolved Hide resolved

ValerianRey reviewed

View reviewed changes

tests/unit/aggregation/test_dual_cone_utils.py Outdated Show resolved Hide resolved

ValerianRey reviewed

View reviewed changes

tests/unit/aggregation/test_dual_cone_utils.py Outdated Show resolved Hide resolved

ValerianRey reviewed

View reviewed changes

src/torchjd/aggregation/_dual_cone_utils.py Outdated Show resolved Hide resolved

ValerianRey reviewed

View reviewed changes

src/torchjd/aggregation/_dual_cone_utils.py Outdated Show resolved Hide resolved

ValerianRey reviewed

View reviewed changes

src/torchjd/aggregation/_gramian_utils.py Outdated Show resolved Hide resolved

ValerianRey reviewed

View reviewed changes

src/torchjd/aggregation/_gramian_utils.py Outdated Show resolved Hide resolved

ValerianRey reviewed

View reviewed changes

src/torchjd/aggregation/_dual_cone_utils.py Outdated Show resolved Hide resolved

ValerianRey reviewed

View reviewed changes

tests/unit/aggregation/test_dual_cone_utils.py Outdated Show resolved Hide resolved

ValerianRey reviewed

View reviewed changes

src/torchjd/aggregation/_dual_cone_utils.py Outdated Show resolved Hide resolved

ValerianRey reviewed

View reviewed changes

src/torchjd/aggregation/_dual_cone_utils.py Outdated Show resolved Hide resolved

ValerianRey reviewed

View reviewed changes

src/torchjd/aggregation/_dual_cone_utils.py Outdated Show resolved Hide resolved

PierreQuinton added 6 commits

February 1, 2025 18:54


          Factorize projection onto the dual cone.

aa9897e


          Move regularization out of projection onto the dual cone.

f534006


          Improve code coverage

864d7a4


          Add link to paper in weights_of_projection

b0739fe


          Rename functions in _dual_cone_utils

cd43f65


          Rename lagrangian to lagrangian_multipliers to match the maths

402024c

ValerianRey reviewed

View reviewed changes

src/torchjd/aggregation/_gramian_utils.py Outdated Show resolved Hide resolved

PierreQuinton and others added 8 commits

February 2, 2025 09:44


          Set explicitly the dtype and device of regularization matrix

fb07035


          Rename functions & variables, add _to_array

756e8c5


          Fix bug of transposed output

601bc6e


          Delete outdated tests

b3615cd


          Rename functions _normalize etc

c9afb41


          Improve docstring and naming of test of kkt conditions.

e141600


          Improve docstring of _get_projection_weights

9c7eaac


          Add test of tensorization of _get_projection_weights

f04581e

PierreQuinton commented

View reviewed changes

tests/unit/aggregation/test_dual_cone_utils.py Show resolved Hide resolved

PierreQuinton commented

View reviewed changes

src/torchjd/aggregation/_dual_cone_utils.py Outdated Show resolved Hide resolved

PierreQuinton added 2 commits

February 3, 2025 09:38


          Clarify the QP problem that is solved, it now matches exactly that of…

a5f562a

… Proposition 1 instead of the dual of projection onto the dual cone.


          Improve docstring of _get_projection_weight_vector

627adf0

Contributor

ValerianRey commented Feb 3, 2025

I think a5f562a is worth a changelog entry, with a vague explanation of what was done and how it may affect results. For instance

Changed

Refactored the underlying optimization problem that UPGrad and DualProj have to solve to project onto the dual cone. This may minimally affect the output of these aggregators.

PierreQuinton added 2 commits

February 4, 2025 08:22


          Add changelog entry

bb411c9


          Add no_grad context to the computation of UPGrad and DualProj as we c…

782aad6

…annot infer any grad from our solve using qp_solvers.

ValerianRey reviewed

View reviewed changes

src/torchjd/aggregation/upgrad.py Outdated Show resolved Hide resolved

ValerianRey added 9 commits

February 4, 2025 20:43


          Improve dual cone utility functions

ffca7d0

- Change parameter order
- Change variable names
- Simplify operations
- Rename functions
- Improve documentation


          Improve variable names:

60a3851

- Always use G for the gramian
- Always use u and U for a vector / tensor of weights prior to projection
- Always use w and W for a vector / tensor of projection weights
- Add underscore to disambiguate with numpy version when necessary


          Improve docstring of test_solution_weights

c7a2fe8


          Remove torch.no_grad in DualProj and UPGrad

633b286


          Add docstring to test_tensorization_shape

4f3e48a


          Improve docstring formatting

0c224fc


          Simplify UPGrad

c759201

- I think it's better to not name self.weighting(matrix) as u, because it doesn't have the same role as a single u afterwards. Therefore, I removed this variable name.


          Add docstring to _regularize

c604c80


          Improve formatting

20aa85c

Contributor

ValerianRey commented Feb 4, 2025 •

edited

Loading

✅ Code: LGTM
✅ Correctness: Tests pass on cpu (in CI) and on gpu (on my machine).
✅ Results: I ran the trajectories on the EWQ and CQF objective with this version, and the results are visually exactly the same as with v0.5.0.
✅ Runtime: The tests on cpu gave very similar runtimes, and the tests on gpu gave slightly better (3 to 6% I would say) runtimes with this version.

ValerianRey changed the title ~~Factorize dual cone projections~~ Refactor dual cone projections

ValerianRey merged commit 1d87f56 into main

14 checks passed

ValerianRey deleted the factorize_dual_cone_projections branch

February 4, 2025 23:17

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

package: aggregation refactor