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

Reducing compilation times #514

Draft
wants to merge 3 commits into
base: main
Choose a base branch
from
Draft

Conversation

vroulet
Copy link
Collaborator

@vroulet vroulet commented Aug 31, 2023

The plan of this PR is to reduce the number of times the objective function is compiled.
I started by adding a test in common_tests to see how many times a function is compiled for each solver, see below for the results.
I will focus first on a simple solver like gradient descent to see if we can reduce the number of compilations.
Then I'll extend to each solver until the aforementioned test can be run to satisfy a single compilation of the objective.

AndersonWrapper                       : objective compiled 3 times
BFGS with zoom linesearch             : objective compiled 8 times
ArmijoSGD                             : objective compiled 5 times
GaussNewton                           : objective compiled 26 times
GradientDescent                       : objective compiled 3 times
LBFGS with zoom linesearch            : objective compiled 8 times
LBFGS with hager-zhang linesearch     : objective compiled 75 times
LBFGS with backtracking linesearch    : objective compiled 5 times
LevenbergMarquardt                    : objective compiled 30 times
NonlinearCG with zoom linesearch      : objective compiled 8 times
PolyakSGD                             : objective compiled 3 times
OptaxSolver                           : objective compiled 3 times
AndersonAcceleration                  : objective compiled 3 times
Broyden with backtracking linesearch  : objective compiled 11 times
Bisection                             : objective compiled 4 times
BlockCoordinateDescent                : objective compiled 5 times
LBFGSB with zoom linesearch           : objective compiled 8 times
ProjectedGradient                     : objective compiled 3 times
ProximalGradient                      : objective compiled 3 times
MirrorDescent                         : objective compiled 2 times
BacktrackingLineSearch                : objective compiled 5 times
HagerZhangLineSearch                  : objective compiled 40 times
ZoomLineSearch                        : objective compiled 6 times

@vroulet
Copy link
Collaborator Author

vroulet commented Sep 4, 2023

I'm now printing the type of the input:

  • If the function is jitted:
    • At compilation time, this type must be a Tracer.
    • Afterwards, nothing will be printed.
  • Otherwise if the function is not jitted,
    • no compilation happens, so the printed type is not Tracer.

The results are given below. The current implementation is not as bad as previously claimed: for e.g. the OptaxSolver there is only one compilation (during the update). The call to the objective in the init function (at least in OptaxSolver/PolyakSGD but I think in other solvers too) does not seem to incur additional compilations.

AndersonWrapper                       : objective compiled 2 times
BFGS with zoom linesearch             : objective compiled 7 times
ArmijoSGD                             : objective compiled 3 times
GaussNewton                           : objective compiled 25 times
GradientDescent                       : objective compiled 2 times
LBFGS with zoom linesearch            : objective compiled 7 times
LBFGS with hager-zhang linesearch     : objective compiled 38 times
LBFGS with backtracking linesearch    : objective compiled 3 times
LevenbergMarquardt                    : objective compiled 29 times
NonlinearCG with zoom linesearch      : objective compiled 7 times
PolyakSGD                             : objective compiled 1 times
OptaxSolver                           : objective compiled 1 times
AndersonAcceleration                  : objective compiled 2 times
Broyden with backtracking linesearch  : objective compiled 6 times
Bisection                             : objective compiled 2 times
BlockCoordinateDescent                : objective compiled 5 times
LBFGSB with zoom linesearch           : objective compiled 7 times
ProjectedGradient                     : objective compiled 2 times
ProximalGradient                      : objective compiled 2 times
MirrorDescent                         : objective compiled 2 times
BacktrackingLineSearch                : objective compiled 3 times
HagerZhangLineSearch                  : objective compiled 40 times
ZoomLineSearch                        : objective compiled 6 times

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant