Tensorflow and PyTorch Setup Guide with Jupyter Notebook and CUDA

Befor we continue we have to make sure what is the latest python version for supported Tensorflow.

For me it is Python 3.9–3.12.

Now we have to install tensorflow and pytorch. We can install all of these using python. It’s better we install python environment so that we don’t pollute our system. There are many ways to create python environment. But we are going to use conda.

We are going to install miniconda and start from there. Go to miniconda website and download.

Make sure to install miniconda in a location where you have enough storage. Our directory will become big overtime.

curl -O https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh
sh Miniconda3-latest-Linux-x86_64.sh

We have to check inf conda is available in terminal.

conda --version

If it’s available create python environment. Since Latest tersorflow supported python version is 3.12 as of now. We create:

conda create -n ml-env python=3.12 -y

Then activate it:

conda activate ml-env

Check where conda env is located:

conda env list

Now install Jupyter Notebook inside conda environment.

conda install notebook -c conda-forge -y

conda-forge is community driven channel. notebook just contain basic jupyterlab. install jupyterlab for full package.

To make conda available inside jupyter notebook. So, that we can install any dependency inside jupter notebook.

# Install the IPython kernel package
conda install ipykernel -y

# Register this environment with Jupyter
python -m ipykernel install --user --name ml-env --display-name "Python (ml-env)"

If you have NVIDIA gpu. Make sure to install CUDA drivers from NVIDIA website. And check out PyTorch website for different version that is available.

CPU

python3 -m pip install -U 'tensorflow[and-cuda]'
python3 -m pip install -U torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu

conda support of PyTorch and Tensorflow has beed deprecated. They now push directly through pip.

NVIDIA GPU

python3 -m pip install -U 'tensorflow[and-cuda]'
python3 -m pip install -U tensorrt tensorrt-bindings tensorrt-libs --extra-index-url https://pypi.nvidia.com
python3 -m pip install -U torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu${CUDA_TK_VER//./}

Now run jupyter notebook

jupyter notebook

Select the Kernel:

In the Jupyter interface, open a notebook (or create a new one).

From the top menu, go to Kernel > Change kernel > Select Python (ml-env) from the list.

Now your notebook will run using the ml-env environment and have access to all the packages you’ve installed there (including PyTorch, TensorFlow, etc.).

Let’s test if Cuda working.

Test from terminal Tensorflow GPU Test

python -c "import tensorflow as tf; device = 'GPU' if tf.config.list_physical_devices('GPU') else 'CPU'; print(f'TensorFlow is running on: {device}')"

PyTorch GPU Test

python -c "import torch; device = 'GPU' if torch.cuda.is_available() else 'CPU'; print(f'PyTorch is running on: {device}')"

Test from Jupyter Notebook Cell For Tensorflow

import tensorflow as tf

print("TensorFlow CUDA Available:", tf.test.is_built_with_cuda())
print("GPU Available:", tf.config.list_physical_devices('GPU'))

For PyTorch

import torch

print("PyTorch CUDA Available:", torch.cuda.is_available())
print("Device Name:", torch.cuda.get_device_name(0) if torch.cuda.is_available() else "No GPU")