Installing Jupyter Lab for ROCm

Set up Jupyter Lab for Deep Learning on a ROCm system

How to Install Jupyter Lab for ROCm

In an earlier post we explained how to deploy the ROCm libraries for PyTorch and the scaffolding code needed to train a basic model on the GPU. We used a simple command line script to validate the software and hardware stack. Now let’s create a Jupyter Lab server running in a Python virtual environment for deep learning work using a ROCm system.

A Virtual Environment for PyTorch and Jupyter Lab

A Python virtual environment is a walled-off runtime system that is separate from the system-wide Python installed in the local host. The first CLI command below creates a new virtual environment, with all its files stored in the .venv folder.

python3 -m venv .venv
source .venv/bin/activate

(.venv) ~$ deactivate

We switch out of the system-wide Python into the virutal environment – activating the VENV – when we execute the second command. The (.venv) prompt is a visual indicator that the virtual environment is active. Typing another CLI command, deactivate, terminates the VENV and returns us to the system-wide Python context.

With the virtual environment active, we proceed to install the PyTorch and Jupyter Lab. Here we install the version 7.2 ROCm runtime system.

(.venv) ~$ python3 -m pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm7.2
(.venv) ~$ python3 -m pip install jupyterlab

At this point Jupyter Lab is set and it comes preloaded with a basic Python kernel (labeled Python 3 ipykernel). Because the default Jupyer Lab kernel runs notebook scripts only on the CPU, we need to take additonal steps to hook the Jupyter Lab server with a Python kernel suitable for ROCm.

(.venv) ~$ python3 -m pip install ipykernel
(.venv) ~$ python3 -m ipykernel install --user --name=rocm_env

(.venv) ~$ jupyter lab    # Server listens to port 8888

Now when we start Jupyter Lab sever, we see two kernels launchers: python3 and rocm_env. We can choose either to run our Jupyter notebook or console script.

Validating the Jupyter Lab and ROCm System

The code below offers a simple test for our Jupyter Lab and ROCm set up. The script adapts Moroney’s code which creates a basic vision recognition model and trains it on the Fashion MNIST dataset. The Python pickle file model_state-best_loss.pth stores the weights and biases of an instance of FashionMNISTModel class. As we see in the script below, we take the pre-trained weights and biases and load it into a new instance of FashionMNISTModel.

The internal dictionary of an instance of the FashionMNISTModel was stored to model_state-best_loss.pth in pickle format. Here is the code for extracting those trained weights and biases.

import os; os.environ["HSA_OVERRIDE_GFX_VERSION"] = "11.0.0"
import torch
import torch.nn as nn


# Redefine the exact same architecture
class FashionMNISTModel(nn.Module):
    def __init__(self):
        super(FashionMNISTModel, self).__init__()
        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
                nn.Linear(28*28, 128),
                nn.ReLU(),
                nn.Linear(128, 10),
                nn.LogSoftmax(dim=1)
        )

    def forward(self, x):
        x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        return logits

# Define the target GPU device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Instantiate the blank model and move it to the GPU
model = FashionMNISTModel().to(device)

# Load the raw weights from the .pth file into the model
# The weights represent the state_dict saved earlier
model.load_state_dict(torch.load("model_state-best_loss.pth", weights_only=True))

# Lock the model into evaluation mode
model.eval()

print("Model successfully loaded and ready for inference on:", device)

The code above is tailored for ROCm. After we create a Jupyter notebook tied to the rocm_env kernel, and copy-paste the Python code, we can proceed to execute the script on the local AMD GPU.

Model successfully loaded and ready for inference on: cuda

Once the message above appears in our Jupyter notebook, we’ve successfully validated the connection between Jupyter Lab and ROCm subsystem.