Inference is really slow?

[florisvb]

Is your feature request related to a problem? Please describe.

After training a model, running model.predict is really slow.. 100-200x slower than I think it should be. Perhaps it is user error but I did not come across a way to make this faster? Or maybe it is a bug if inference is supposed to be faster?

Inference for a very simple model takes ~5 ms, but the underlying numpy computations should only take 20 µs.

Here is a speed test result:
https://github.com/florisvb/Nonlinear_and_Data_Driven_Estimation/blob/main/Lesson_15_SINDY/speed_test.png

Describe the solution you'd like

Two options I can think of:

1. Modify the predict method to be faster.. ie. is there stuff that predict is computing that is unnecessary? Is there an issue with the types? Is there a way to get a sparse model so that we're not wasting computer on features that will be zeroed out anyways?
2. Export the model as a fast inference model. For example, can the non zero features and their coefficients be exported to a portable file type and be reconstructed as a super fast function? This would be really nice to have, regardless.

Describe alternatives you've considered

I could do option 2 with enough time, but it would be nice if there were a pysindy method to save and load a model? Or maybe there is and I did not find it.

Additional context

This is not a minimal example, but it is a fully functional notebook that shows the issue. It will run in google colab, but takes a little time to install everything and download data used in pysindy:
https://github.com/florisvb/Nonlinear_and_Data_Driven_Estimation/blob/main/Lesson_15_SINDY/A_planar_drone_pysindy_speed_test.ipynb

[florisvb]

I don't know pysindy super well, but for my scenario I was able to build an inference model that improves the speed by 100x. Maybe with a little fender and body work for more general libraries this could provide a template for a general compilation?

My feature names:

['1',
 'theta_dot',
 'x',
 'x_dot',
 'z',
 'z_dot',
 'j1',
 'j2',
 'sin(theta)',
 'cos(theta)',
 'j1*sin(theta)',
 'j1*cos(theta)',
 'j2*sin(theta)',
 'j2*cos(theta)']

Inference model:

class SINDY_INFERENCE_MODEL:
  def __init__(self, model, symbols):
    self.feature_names = copy.copy(model.get_feature_names())
    self.coefficients = copy.copy(model.coefficients())
    self.symbols = symbols

    self.lambda_funcs = self.build_lambda_funcs(self.feature_names, self.symbols)
  
  def build_lambda_funcs(self, feature_names, symbols):
    # replace ' ' with multiplication
    for i, feature in enumerate(feature_names):
      feature = feature.replace(' ', '*')
      feature_names[i] = feature
    
    # Parse feature names into SymPy expressions
    sympy_exprs = []
    for name in feature_names:
        expr = sp.sympify(name.replace('^', '**'))
        sympy_exprs.append(expr)

    # Convert to fast lambda functions
    lambda_funcs = [sp.lambdify(symbols, expr, 'numpy') for expr in sympy_exprs]

    return lambda_funcs 
  
  def predict(self, x, u):
    input = np.hstack([np.ravel(x), np.ravel(u)])
    results = np.array([func(*cols) for func in lambda_funcs])
    v = self.coefficients@np.atleast_2d(results).T
    return v

Usage

# I think this is necessary, also any other functions in the library
# I only have sin and cos in my library
sin = sp.sin
cos = sp.cos 

symbols = [theta, theta_dot, x, x_dot, z, z_dot, j1, j2]
sindy_inference = SINDY_INFERENCE_MODEL(model, symbols).predict

x_test = np.atleast_2d(np.random.randn(1, 6))
u_test = np.atleast_2d(np.random.randn(1, 2))
%timeit sindy_inference(x_test, u_test)

This yields about 50 micro sec per call, improved from 5 ms.

This could probably be improved even further by only evaluating lambda funcs for coefficients that are nonzero. But this is fast enough that I can use it.

[florisvb]

There were some bugs in that code since I was prototyping in a notebook and forgot to run from the start.

But here is a complete and functional notebook (using the generic pysindy example) that does two things:

1. Builds a fast inference model from lambda functions (100x faster than model.predict for individual inference calls)
2. Saves the model as a json file that is very portable across systems

https://github.com/florisvb/Nonlinear_and_Data_Driven_Estimation/blob/main/Lesson_15_SINDY/pysindy_fast_inference.ipynb

I have only tested it with the basic pysindy example, and my use case which has sin and cos in the library. But it may work for broader libraries.

I should note that Claude AI wrote almost all of this code at my direction. Thanks to good pysindy docs I think it handled it remarkably well!!

[Jacob-Stevens-Haas]

Thanks! This is definitely time (no pun intended) to add airspeed velocity timing benchmarks to pysindy, so that this stuff gets tracked in CI. Your speed test makes a lot of sense.

My guess is that the speed is due to using a scikit-learn Pipeline. We're trying to remove this, so I'd say the first step is probably adding benchmark, removing Pipeline, and then profiling the code if that does't speed it up. And only then, if necessary, figuring a way to only evaluate features with nonzero coefficients.

I'd say we're also looking to make more code general to cvxpy, jax, and numpy, and choosing which array library to use depending on input types. If that's the case, jax's JIT might come in useful. There's also been a proposal to allow selecting a subset of a feature library, and its possible that all the libraries could implement a getitem to do something special.

Models should be able to save and load with pickle. I believe there's a test to ensure all PRs preseve the pickle-ability. If it doesn't unpickle correctly, it's a bug