Nx / Axon - CUDA support


I’ve just started learning Elixir few weeks ago and I really enjoy that language.
Currently, I’d like to start with machine learning / deep learning using Nx and Axon. I have an old NVIDIA GPU (NVS 510). In Python I could use PyTorch with CUDA, but for Elixir there aren’t that much tutorials showing how to set everything up. I’ve seen some articles about Axon and in all of them they use EXLA.
How to use full potential of my computer? Should I do everything just as it is shown in articles or is there a better way?

Have a look at this post https://elixirforum.com/t/compiling-xla-with-cuda-11-2-support-needed/42528/4.

I usually have to set these env variables XLA_TARGET=cuda111 EXLA_TARGET=cuda LIBTORCH_TARGET=cuda

At the moment looks like Cuda 11.1 is prebuilt which might not be supported. From the post linekd above, you might have to do a custom Bazel build, but I’m not sure

1 Like

Thanks! I’ll definitely try that.
If I understood correctly, even EXLA on CPU should run faster than training on CPU / GPU?

Yes should be faster than plain Elixir on CPU. Using GPU would be ideal, but for smaller datasets, XLA CPU is pretty fast.

You can find some stats here (just find the word benchmarks)