search

knights-analytics / hugot

Huggingface transformer pipelines in Golang
Apache License 2.0
279 stars 13 forks source link

GPU support #9

Closed riccardopinosio closed 6 months ago

riccardopinosio commented 7 months ago

Since we rely on onnxruntime it should be fairly straightforward to support GPU for the pipelines.

sedletsky-f5 commented 6 months ago

is there a plan to support GPU?

riccardopinosio commented 6 months ago

@sedletsky-f5 hi, there sure is, as onnxruntime should make that relatively painless. I haven't gotten around to it yet, wanted to add support for a couple more pipelines first, but could look at it earlier if you think it's important for you application

RJKeevil commented 6 months ago

Today i started a new branch (gpu-support) which exposes all the options for the available providers. If you want to try a pre alpha version i think the remaining work is to select an appropriate gpu-enabled onnxruntime library from https://github.com/microsoft/onnxruntime/releases , then test this in an environment with the right OS dependencies/libraries. Let us know if you try this early.

RJKeevil commented 6 months ago

Cuda is now fully working for me on the GPU branch, so now would be a good time to try it if you like. We'll release this imminently.

sedletsky-f5 commented 6 months ago

sorry, missed previous messages - yes, it is important for my project and I will test gpu-support branch shortly. thanks a lot!

riccardopinosio commented 6 months ago

Hi @sedletsky-f5, this is now implemented in main (new release 0.1.0).

sedletsky-f5 commented 6 months ago

great news, @riccardopinosio ! I've exactly tested it on my env and GPU version works properly and performance is also much better than CPU version (about 5 times faster for running feature extraction pipeline).

now, I have some improvement suggestion - today users should decide what version to use and even specify the selected library with WithOnnxLibraryPath(), but it can be done automatically using function gotch.CudaIfAvailable() available in https://github.com/sugarme/gotch/