Using GPU to enhance the speed

[JamesWeiChen]

@frankmj Since HSSM is built on PyMC and incorporates features from the Python Bayesian ecosystem, it suggests it can utilize GPU acceleration indirectly through PyMC if appropriately configured. However, I think we cannot currently use GPU to boost the speed, right? Is it possible to have that feature in the future? Thank you!

Best,
James

[frankmj]

Yes you can certainly use GPUs and this should speed up processing considerably especially when using LANs. If using Colab you can simply switch the runtime to use GPU. Michael

…

On Tue, Apr 23, 2024 at 10:19 PM James ***@***.***> wrote: @frankmj <https://github.com/frankmj> Since HSSM is built on PyMC and incorporates features from the Python Bayesian ecosystem, it suggests it can utilize GPU acceleration indirectly through PyMC if appropriately configured. However, I think we cannot currently use GPU to boost the speed, right? Is it possible to have that feature in the future? Thank you! Best, James — Reply to this email directly, view it on GitHub <#397>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAG7TFE5MO3B65JZE7P7JMTY64JDDAVCNFSM6AAAAABGWAXKZSVHI2DSMVQWIX3LMV43ERDJONRXK43TNFXW4OZWGU2TCNJXGY> . You are receiving this because you were mentioned.Message ID: ***@***.***>

[JamesWeiChen]

@frankmj Thank you for your suggestion! Could you provide specific instructions on configuring and enabling GPU acceleration on a personal computer? If there are specific commands or steps, could you guide me through them? Thank you!

[frankmj]

I will let @paul Xu ***@***.***> respond on this issue :^)

…

On Wed, Apr 24, 2024 at 10:27 AM James ***@***.***> wrote: @frankmj <https://github.com/frankmj> Thank you for your suggestion! Could you provide specific instructions on configuring and enabling GPU acceleration on a personal computer? If there are specific commands or steps, could you guide me through them? Thank you! — Reply to this email directly, view it on GitHub <#397 (reply in thread)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAG7TFGJH75N7JPVQJ6WV6TY666MFAVCNFSM6AAAAABGWAXKZSVHI2DSMVQWIX3LMV43SRDJONRXK43TNFXW4Q3PNVWWK3TUHM4TEMJUGI3TQ> . You are receiving this because you were mentioned.Message ID: ***@***.***>

[digicosmos86]

Hi @JamesWeiChen,

It's actually quite easy to use GPU acceleration with HSSM. All you need to do is to install JAX with GPU support. Detailed installation guide can be found here. Once you have installed JAX, proceed to install pymc and HSSM following our instruction here.

Once everything is set up, JAX will automatically perform computation with GPU unless otherwise instructed. You just need to use a JAX-based sampler (e.g. set sampler="nuts_numpyro" when calling model.sample()).

If you are using colab with GPU, JAX with GPU support is installed by default. You can just install HSSM via pip. Please check out this issue to get started.

Please let me know if you have any questions :)

Paul