tensorflow , pytorch + ubuntu focal 20.04 + conda

I thought it’s trivial in the first place, but in the end, it’s not, but still simple if all parties do it right on compatibility!

There are so many combinations of NV card, driver, installation method, etc to make it maybe hopeless to just install. Well, then maybe it makes sense to have a place to just run and keep making sure it works. I believe it’s tested somewhere, maybe just not have one single place to have those all.

So, here it is:

Nvidia driver:

$ sudo apt-get install nvidia-driver-460 — no-install-recommends
$ sudo apt-get install nvidia-utils-460-server # to get nvidia-smi

For today, above install 460.56. Some previous versions might not work. I believe most new NV GPU can be supported by this one.

Conda: use python 3.8 miniconda

command to install PyTorch and TensorFlow

$ conda install pytorch torchvision torchaudio \
cudatoolkit=10.1 \
matplotlib jupyterlab ipywidgets \
tensorflow-gpu \
-c pytorch

That’s it.

Let me know if the above one does not work, I’ll try to see if I can do a test.

Some background:

PyTorch has a conda channel and supports many CUDA versions. For example, they have versions for cuda10 and cuda11 separately. For tensorflow, I check several conda channels but I simply can’t find any Tensorflow-gpu conda that consume cuda 11.