search

DrChainsaw / ONNXNaiveNASflux.jl

Import/export ONNX models
MIT License
43 stars 2 forks source link

Running in Cuda or AMD? #86

Closed jdiaz97 closed 7 months ago

jdiaz97 commented 11 months ago

Hello! The docs are not clear about GPU support (maybe it's too implied? I'm not a fluent Flux user) I know Flux can support AMD in certain scenarios, could I run onnx models using this and AMD?

DrChainsaw commented 11 months ago

Yes that should be possible. The docs are a bit sparse in an attempt to not repeat stuff already found in the docs for NaiveNASflux.

The tl;dr is that a loaded model is fully Functor compatible, so things like model_gpu = load("model.onnx") |> gpu should just work.

I haven't tested this with an AMD gpu since I only have a Nvidia, but everything that works for a Flux.Chain w.r.t parameter movement should work for models loaded with this package.