How to make Jupyter Notebook to run on GPU?

How to make Jupyter Notebook to run on GPU?

Jupyter Notebook is one of the most popular IDEs for data science in python. If you already have Python 3.x and Anaconda installed, you can launch Jupyter Notebook from Anaconda Navigator.

After launching Jupyter Notebook, if you click on New, you will see a dropdown menu to select a virtual environment to choose to launch the notebook. By default, you will only see Python3 environment. Once you select that new notebook page will be opened in separate tab, where you can start doing your coding. This Python3 virtual environment will use your computer’s CPU to compute your code.

It should be enough if you are doing EDA tasks or running simple ML tasks on limited dataset. But if you want to run some heavy computing tasks like training AI and DL models, fine tuning a BERT model, Image processing etc. then you might need an environment with GPU or TPU.

If you have already worked in Kaggle or Google Colab, then you must have seen that there is option in both the platform to choose whether you want to run your model on CPU, GPU or TPU and you can select that runtime env. before launching your notebook. You can do the same thing in your local system if your local system already has a GPU (preferably Nvidia GPU).

Before following below steps make sure that below pre-requisites are in place:

  1. Python 3.x is installed.
  2. Anaconda is installed.
  3. CUDA toolkit is installed.

Steps to run Jupyter Notebook on GPU

  1. Create a new environment using conda:

Open command prompt with Admin privilege and run below command to create a new environment with name gpu2.

conda create -n gpu2 python=3.6

Follow the on-screen instructions as shown below and gpu2 environment will be created.

  • Run below command to list all available environments.

conda info -e

  • Now, run below command to activate / enable newly created gpu2 environment.

conda activate -n gpu2

  • Install tensorflow-gpu. Here I have installed tensorflow-gpu v2.3.0.

You can check below link to find the compatibale tensorflow-gpu version with your install Python version.

pip install tensorflow-gpu==2.3.0

  • Use tf.test.is_built_with_cuda() to validate if TensorFlow was built with CUDA support. You can see below it’s returning True.
  • Install ipykernal by running below command. Before running this make sure that you already have activated gpu2 environment (step 3).

conda install -c anaconda ipykernel

  • Now install the new kernel by running below command:

python -m ipykernel install –user –name=gpu2

  • Now, this new environment (gpu2) will be added into your Jupyter Notebook. Launch Jupyter Notebook and you will be able to select this new environment.
  • Launch a new notebook using gpu2 environment and run below script. It will show you all details about the available GPU.

CUDA support is also available.

Any notebook created in gpu2 environment will use the GPU to compute and if you need only CPU to compute then you can launch the notebook in Python3 environment.

Anik is an IT professional and Data Science Enthusiast. He loves to spend a lot of time testing and reviewing the latest gadgets and software. He likes all things tech and his passion for smartphones is only matched by his passion for Sci-Fi TV Series.