AMD, ROCM, PyTorch, and AI on Ubuntu: The Rules of the Jungle

Are you the awkward lanky kid trying to join the other cool athletes on the field? Or perhaps you went out to buy an NVIDIA video card, only to realize the 2GB VRAM won’t get you very far? And better yet “upgraded” to an AMD? And then when you got it home and then started hacking away on spacy and Stable Diffusion you realize that it’s a nightmare?

AMD ≠ Nvidia

Environment Variables

Command Line Utilities


NOTE: Multiversion install is not available for the AMDGPU stack.
ainst='sudo amdgpu-install -y --rocmrelease=${ROCMV} --usecase=rocm,rocmdevtools,lrt,hip,hiplibsdk,mllib,mlsdk,dkms'


$ rocminfo | grep 'Name:'
Name: Intel(R) Core(TM) i7-7700K CPU @ 4.20GHz
Marketing Name: Intel(R) Core(TM) i7-7700K CPU @ 4.20GHz
Vendor Name: CPU
Name: gfx803
Marketing Name: Radeon RX 580 Series
Vendor Name: AMD
Name: amdgcn-amd-amdhsa--gfx803

Now, for the Common AI Libraries…

$ # oh nose! I need to reinstall spacy!
$ pip uninstall torch torchaudio torchvision spacy -y
$ pip install torch torchvision torchaudio --extra-index-url --no-cache-dir
$ pip install cupy # or install from source
$ pip install spacy
  • ROCM
  • torch / cupy
  • libraries that rely on touch / cupy

End of the Quick Walk-through



Senior Full Stack Developer & Tech Lead (#openforwork)

Get the Medium app

A button that says 'Download on the App Store', and if clicked it will lead you to the iOS App store
A button that says 'Get it on, Google Play', and if clicked it will lead you to the Google Play store