MLJ.jl icon indicating copy to clipboard operation
MLJ.jl copied to clipboard

Restore machine and inference on GPU

Open Moelf opened this issue 1 year ago • 4 comments

After saving and loading a machine, it runs inference on GPU (specifically, for XGBoost), is there an interface to move the model back to GPU?

Moelf avatar Nov 23 '24 06:11 Moelf

I'm not sure I understood your question. Could you provide example code.

OkonSamuel avatar Dec 07 '24 20:12 OkonSamuel

if you train a machine, and then:

mach = restore!(serializable(mach_bdt))

any subsequent call to predict(mach, X) would run on CPU not GPU. In this particular case the wrapper for XGBoost.jl is lacking this capability, but in general, does MLJ even have an interface for "moving model onto GPU"?

Moelf avatar Dec 07 '24 20:12 Moelf

Oh okay, I see what you mean now. At the moment we don't have support for gpu training with XGboost model. But I believe that since XGBoost.jl supports this, it should be possible to add support for this through MLJ. One way would be to add a method to the interface that gets overloaded by the different models. At the moment for MLJFlux models this is done in a locally but it may be more beneficial to generalize this.

OkonSamuel avatar Dec 07 '24 20:12 OkonSamuel

GPU training for XGBoost model is okay: https://github.com/JuliaAI/MLJXGBoostInterface.jl/issues/54

I think the problem is more that there's no way to move inference back to GPU -- I think in thsi case XGBoost.jl may need the DMatrix to live on GPU or something for this to work

Moelf avatar Dec 07 '24 23:12 Moelf