Gradient Computation/Assembly Very Slow

Question

I'm implementing an existing MATLAB optimization code in Python using OpenMDAO.

The basic problem structure from openmdao summary is:

============== Problem Summary ============
Groups:              14
Components:          70
Max tree depth:       5

Design variables:            7   Total size:      121

Nonlinear Constraints:      12   Total size:      936
    equality:                0                      0
    inequality:             12                    936

Linear Constraints:          0   Total size:        0
    equality:                0                      0
    inequality:              0                      0

Objectives:                  1   Total size:        1

Input variables:           357   Total size:    41451
Output variables:          271   Total size:    31855

Total connections: 357   Total transfer data size: 41451

At the top level, there are 5-10 components for evaluating constraints and objectives, there is a group where the DOF variables are being defined/parsed (SplineComp) and a large group where the bulk of the computations are being done (contains one group with two levels, one group with one level, and a bunch of components). This is basically the same structure as the MATLAB code. On the MATLAB side, I am computing analytical gradients at the problem level, but calling the large group 14 times for finite difference gradients of that portion. In Python, I'm using a mix of hand calculated gradients and automatic differentiation tools of the components inside that large group and its subgroups.

I've verified that the results match between Python and MATLAB for a prescribed set of inputs and check_partials/check_totals give satisfactory results in Python, I'm now stumped at the speed comparisons in benchmarking.

In MATLAB, the gradient computation takes ~14x longer than a standalone evaluation of the objective/constraint, which is as expected given the finite difference approach used.

In Python, the compute_totals takes ~25-50x longer than run_model. This is perplexing to me as I'd expect much better performance since there's zero finite difference involved here. I'm not expecting compute_totals to be the same speed as run_model, but it's more than an order of magnitude slower than I expected.

I ran the command line profiling (iprof) and unsurprisingly most of the time is being spent in solve_linear, but after using the python profile tool, it's showing nearly half the time is spent in dictionary_jacobian.

I tried looking into using the coloring feature to improve the computation time, but it doesn't seem to be finding the actual sparsity of the jacobian (reports ~64% non-zero compared to ~6% actual):

I've verified that all my custom components are using sparse partial definitions through the rows/cols option of declare_partials (do certain built in components not use sparse definitions that would cause this? I'm using a few inside of the large group for convenience) I know the hackathon findings (https://github.com/OpenMDAO/RevHack2020/blob/master/solution_approaches/how_big.md) suggest I should have one large component instead of multiple small components, but that'll be hard to implement here without losing a good chunk of the flexibility I need for activating different features inside that group.

Justin Gray Justin Gray · Accepted Answer · 2020-11-09T20:06:00

If you need the flexibility then keep the components the size you have them.

Dictionary Jacobian implies that you're not taking advantage of the DirectSolver which is needed to get really good performance for models of your size. This includes the use of the assembled Jacobian, which is also needed for really fast derivatives.

Try setting the DirectSolver to the linear solver at the top of your model like this:

prob.model.linear_solver = om.DirectSolver()

For generality, I should not the following:

You can only use DirectSolver on a serial group (but that serial group could be part of a larger distributed model).
You can use a DirectSolver on a sub-group of a model. Depending on model structure, this may be faster (See section 5.3 of this paper)
DirectSolver works with an assembled jacobian (as opposed to the DictionaryJacobian referenced in question). This is down in the weeds, but in the vast majority of cases you want a sparse assembled jacobian for the direct solver. Very occasionally a dense assembly is better. That is controlled through system option assembled_jac_type which defaults to csc (a sparse format). You might try dense as a test -- but again, csc is very likely to be fastest.

Gradient Computation/Assembly Very Slow

1 Answers