Configure a Jupyter notebook to use GPUs for AI/ML modeling

For the best experience in this learning path, we suggest that you complete the following learning resources in the order shown. When you click on a resource, it will open in a new tab. Keep this page open so you can easily move on to the next resource!

    

In the previous section, we have:

  1. Imported torch libraries (utilities).
  2. Listed available GPUs.
  3. Checked that GPUs are enabled.
  4. Assigned a GPU device and retrieved the GPU name.
  5. Loaded vectors, matrices, and data onto a GPU.
  6. Loaded a neural network model onto a GPU.
  7. Trained the neural network model.

Let’s now determine how our simple torch model performs using GPU resources.  

In the getting-started-with-gpus directory, double click on the torch-test-model.ipynb file (highlighted in Figure 9) to open the notebook.

Figure 9:  Double-click on the torch-test-model.ipynb file to open the notebook.
Figure 9: Double-click on the torch-test-model.ipynb file to open the notebook.

 

After importing the torch and torchvision utilities, assign the first GPU to your device variable.  Prepare to import your trained model, then place the model on your GPU and load in its trained weights:

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import datasets
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(device)  # let's see what device we got
# Getting set to import our trained model.

# batch size of 1 so we can look at one image at time.
batch_size = 1


class SimpleNet(nn.Module):
    def __init__(self):
        super(SimpleNet, self).__init__()
        self.fc1 = nn.Linear(784, 784)
        self.fc2 = nn.Linear(784, 10)

    def forward(self, x):
        x = x.view(batch_size, -1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.fc2(x)
        output = F.softmax(x, dim=1)
        return output
model = SimpleNet().to( device )
model.load_state_dict( torch.load("mnist_fashion_SimpleNet.pt") )

You are now ready to examine some data and determine how your model performs.  The sample run in Figure 10 shows that the model predicted a "bag" with a confidence of about 0.9192. Despite the % in the output, 0.9192 is very good because a perfect confidence would be  1.0.

 

Figure 10: The model classifies a picture as a bag, with a confidence of about 0.9192.
Figure 10: The model classifies a picture as a bag, with a confidence of about 0.9192.

 

This concludes our tutorial on how to get started with NVIDIA GPUs.

Previous resource
How to examine GPU resources with PyTorch