How to use GPU on Jupyter notebook in Google Cloud AI Platform

Question

I'm using a GPU-enabled Jupyter notebook on the Google Cloud AI Platform.

However, despite being GPU enabled, it is only using the CPU when I train my machine learning model.

Does anyone know how I can make it use GPU instead? I've been using Google Colab notebooks before this, and there you can simply change the runtime to use GPU in the menu at the top of the notebook. But it doesn't look like there is an option like that in the Jupyter notebook.

Appreciate any advice on how to get it to use GPU!

Thanks.

Answer 1

AI Platform Notebooks don't have the drivers pre-installed. When you create an AI Platform Notebooks instance , if you choose to include a GPU, you must select the option to Install NVIDIA GPU driver automatically for me so the image is provisioned the latest stable driver based on the framework's CUDA version.

Answer 2

According to the GCP official documentation you can attach GPU on the Jupyter Notebook. First of all, you need to stop the notebook instance then select the GPU's from the GPU column along with the number of accelerators. After AI Platform Notebooks has finished updating the GPUs for your notebook instance, you can start the instance by selecting your instance and then clicking Start.