How to install pytorch for Ryzen CPU + Nvidia GPU for best performance?

  • CPU - Ryzen 9 3900x
  • GPU - RTX 2080ti
  • OS - Ubuntu 20.04
  • PyTorch - From source

I know this question is asked a lot but I was not able to come to a solution. It is getting confusing for me. So I am asking this on the forums. Can someone comment on this, and point the mistakes I made, or things I missed?

GPU

  1. Install driver (450.), CUDA (11.0), CUDNN (8.0.1).
  2. Add path to .bashrc
  3. Environment variables for pytorch install USE_CUDA=1 USE_CUDNN=1 TORCH_CUDA_ARCH_LIST="7.0"

I think it is enough for GPU. For my use case, I will only train models (no inference setup).

Build options

I found the following environment variables to be useful:

  • USE_FBGEMM=1
  • USE_MKLDNN=0
  • USE_DISTRIBUTED=0 (Reasoning I have only one GPU)
  • USE_PROF=1

CPU libraries

This is where it gets confusing for me, which libraries to install, how to install them.

  • BLAS/OpenBLAS -> OpenBLAS is optimized version of BLAS, so I should use this.
  • LAPACK -> It is built on top of BLAS. So I use OpenBLAS, then it would use OpenBLAS.

Should I compile these libraries from source or use available ubuntu packages? Also, I think cmake will automatically find these libraries so I don’t have to provide any specific build options.

Is this enough for my use case of PyTorch or I missed something? My use case is mostly training models (including mixed precision, quantization).

Non pytorch libraries

This is a random question. If I have installed OpenBLAS on my system, will doing pip install numpy {or some other library} automatically use my system libraries, or will I have to compile these libraries from source also. Assuming I am working in a conda environment.