Home » Tutorials » PyTorch » Quick Start

Quick Start

%matplotlib inline
!pip install ipywidgets
WARNING: The directory '/home/ubuntu/.cache/pip/http' or its parent directory is not owned by the current user and the cache has been disabled. Please check the permissions and owner of that directory. If executing pip with sudo, you may want sudo's -H flag.
WARNING: The directory '/home/ubuntu/.cache/pip' or its parent directory is not owned by the current user and caching wheels has been disabled. check the permissions and owner of that directory. If executing pip with sudo, you may want sudo's -H flag.
Requirement already satisfied: ipywidgets in /usr/local/lib/python3.8/dist-packages (8.0.2)
Requirement already satisfied: ipython>=6.1.0 in /usr/local/lib/python3.8/dist-packages (from ipywidgets) (8.4.0)
Requirement already satisfied: widgetsnbextension~=4.0 in /usr/local/lib/python3.8/dist-packages (from ipywidgets) (4.0.3)
Requirement already satisfied: ipykernel>=4.5.1 in /usr/local/lib/python3.8/dist-packages (from ipywidgets) (6.15.1)
Requirement already satisfied: traitlets>=4.3.1 in /usr/local/lib/python3.8/dist-packages (from ipywidgets) (5.3.0)
Requirement already satisfied: jupyterlab-widgets~=3.0 in /usr/local/lib/python3.8/dist-packages (from ipywidgets) (3.0.3)
Requirement already satisfied: backcall in /usr/local/lib/python3.8/dist-packages (from ipython>=6.1.0->ipywidgets) (0.2.0)
Requirement already satisfied: matplotlib-inline in /usr/local/lib/python3.8/dist-packages (from ipython>=6.1.0->ipywidgets) (0.1.3)
Requirement already satisfied: decorator in /usr/local/lib/python3.8/dist-packages (from ipython>=6.1.0->ipywidgets) (5.1.1)
Requirement already satisfied: pygments>=2.4.0 in /usr/local/lib/python3.8/dist-packages (from ipython>=6.1.0->ipywidgets) (2.13.0)
Requirement already satisfied: stack-data in /usr/local/lib/python3.8/dist-packages (from ipython>=6.1.0->ipywidgets) (0.4.0)
Requirement already satisfied: pexpect>4.3; sys_platform != "win32" in /usr/lib/python3/dist-packages (from ipython>=6.1.0->ipywidgets) (4.6.0)
Requirement already satisfied: setuptools>=18.5 in /usr/local/lib/python3.8/dist-packages (from ipython>=6.1.0->ipywidgets) (41.0.0)
Requirement already satisfied: pickleshare in /usr/local/lib/python3.8/dist-packages (from ipython>=6.1.0->ipywidgets) (0.7.5)
Requirement already satisfied: jedi>=0.16 in /usr/local/lib/python3.8/dist-packages (from ipython>=6.1.0->ipywidgets) (0.18.1)
Requirement already satisfied: prompt-toolkit!=3.0.0,!=3.0.1,<3.1.0,>=2.0.0 in /usr/local/lib/python3.8/dist-packages (from ipython>=6.1.0->ipywidgets) (3.0.30)
Requirement already satisfied: jupyter-client>=6.1.12 in /usr/local/lib/python3.8/dist-packages (from ipykernel>=4.5.1->ipywidgets) (7.3.4)
Requirement already satisfied: psutil in /usr/local/lib/python3.8/dist-packages (from ipykernel>=4.5.1->ipywidgets) (5.9.1)
Requirement already satisfied: tornado>=6.1 in /usr/local/lib/python3.8/dist-packages (from ipykernel>=4.5.1->ipywidgets) (6.2)
Requirement already satisfied: packaging in /usr/local/lib/python3.8/dist-packages (from ipykernel>=4.5.1->ipywidgets) (21.3)
Requirement already satisfied: pyzmq>=17 in /usr/local/lib/python3.8/dist-packages (from ipykernel>=4.5.1->ipywidgets) (23.2.1)
Requirement already satisfied: debugpy>=1.0 in /usr/local/lib/python3.8/dist-packages (from ipykernel>=4.5.1->ipywidgets) (1.6.2)
Requirement already satisfied: nest-asyncio in /usr/local/lib/python3.8/dist-packages (from ipykernel>=4.5.1->ipywidgets) (1.5.5)
Requirement already satisfied: asttokens in /usr/local/lib/python3.8/dist-packages (from stack-data->ipython>=6.1.0->ipywidgets) (2.0.8)
Requirement already satisfied: pure-eval in /usr/local/lib/python3.8/dist-packages (from stack-data->ipython>=6.1.0->ipywidgets) (0.2.2)
Requirement already satisfied: executing in /usr/local/lib/python3.8/dist-packages (from stack-data->ipython>=6.1.0->ipywidgets) (0.10.0)
Requirement already satisfied: parso<0.9.0,>=0.8.0 in /usr/local/lib/python3.8/dist-packages (from jedi>=0.16->ipython>=6.1.0->ipywidgets) (0.8.3)
Requirement already satisfied: wcwidth in /usr/local/lib/python3.8/dist-packages (from prompt-toolkit!=3.0.0,!=3.0.1,<3.1.0,>=2.0.0->ipython>=6.1.0->ipywidgets) (0.2.5)
Requirement already satisfied: jupyter-core>=4.9.2 in /usr/local/lib/python3.8/dist-packages (from jupyter-client>=6.1.12->ipykernel>=4.5.1->ipywidgets) (4.11.1)
Requirement already satisfied: entrypoints in /usr/lib/python3/dist-packages (from jupyter-client>=6.1.12->ipykernel>=4.5.1->ipywidgets) (0.3)
Requirement already satisfied: python-dateutil>=2.8.2 in /usr/local/lib/python3.8/dist-packages (from jupyter-client>=6.1.12->ipykernel>=4.5.1->ipywidgets) (2.8.2)
Requirement already satisfied: pyparsing!=3.0.5,>=2.0.2 in /usr/local/lib/python3.8/dist-packages (from packaging->ipykernel>=4.5.1->ipywidgets) (3.0.9)
Requirement already satisfied: six in /usr/lib/python3/dist-packages (from asttokens->stack-data->ipython>=6.1.0->ipywidgets) (1.14.0)
WARNING: You are using pip version 19.3.1; however, version 22.3.1 is available.
You should consider upgrading via the 'pip install --upgrade pip' command.

This section runs through the API for common tasks in machine learning. Refer to the links in each section to dive deeper.

Working with data

PyTorch has two primitives to work with datatorch.utils.data.DataLoader and torch.utils.data.DatasetDataset stores the samples and their corresponding labels, and DataLoader wraps an iterable around the Dataset.

import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor, Lambda, Compose
import matplotlib.pyplot as plt

Enable Habana

Let’s enable a single Gaudi device by loading the Habana PyTorch plugin library:

#from habana_frameworks.torch.utils.library_loader import load_habana_module
#load_habana_module()
import habana_frameworks.torch.core as htcore

PyTorch offers domain-specific libraries such as TorchTextTorchVision, and TorchAudio, all of which include datasets. For this tutorial, we will be using a TorchVision dataset.

The torchvision.datasets module contains Dataset objects for many real-world vision data like CIFAR, COCO (full list here). In this tutorial, we use the FashionMNIST dataset. Every TorchVision Dataset includes two arguments: transform and target_transform to modify the samples and labels respectively.

# Download training data from open datasets.
training_data = datasets.FashionMNIST(
    root="data",
    train=True,
    download=True,
    transform=ToTensor(),
)

# Download test data from open datasets.
test_data = datasets.FashionMNIST(
    root="data",
    train=False,
    download=True,
    transform=ToTensor(),
)

We pass the Dataset as an argument to DataLoader. This wraps an iterable over our dataset, and supports automatic batching, sampling, shuffling and multiprocess data loading. Here we define a batch size of 64, i.e. each element in the dataloader iterable will return a batch of 64 features and labels.

batch_size = 64

# Create data loaders.
train_dataloader = DataLoader(training_data, batch_size=batch_size)
test_dataloader = DataLoader(test_data, batch_size=batch_size)

for X, y in test_dataloader:
    print("Shape of X [N, C, H, W]: ", X.shape)
    print("Shape of y: ", y.shape, y.dtype)
    break
Shape of X [N, C, H, W]: torch.Size([64, 1, 28, 28]) Shape of y: torch.Size([64]) torch.int64
Code language: CSS (css)

Read more about loading data in PyTorch.

Creating Models

To define a neural network in PyTorch, we create a class that inherits from nn.Module. We define the layers of the network in the __init__ function and specify how data will pass through the network in the forward function. To accelerate operations in the neural network, we move it to Gaudi.

# Use hpu device for training.
# device = "hpu"
device = torch.device("hpu")
print("Using {} device".format(device))

# Define model
class NeuralNetwork(nn.Module):
    def __init__(self):
        super(NeuralNetwork, self).__init__()
        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(28*28, 512),
            nn.ReLU(),
            nn.Linear(512, 512),
            nn.ReLU(),
            nn.Linear(512, 10)
        )

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

model = NeuralNetwork().to(device)
print(model)
Using hpu device NeuralNetwork( (flatten): Flatten(start_dim=1, end_dim=-1) (linear_relu_stack): Sequential( (0): Linear(in_features=784, out_features=512, bias=True) (1): ReLU() (2): Linear(in_features=512, out_features=512, bias=True) (3): ReLU() (4): Linear(in_features=512, out_features=10, bias=True) ) )
Code language: PHP (php)

Read more about building neural networks in PyTorch _.

Optimizing the Model Parameters

To train a model, we need a loss function and an optimizer.

loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)

In a single training loop, the model makes predictions on the training dataset (fed to it in batches), and backpropagates the prediction error to adjust the model’s parameters.

def train(dataloader, model, loss_fn, optimizer):
    size = len(dataloader.dataset)
    model.train()
    for batch, (X, y) in enumerate(dataloader):
        X, y = X.to(device), y.to(device)

        # Compute prediction error
        pred = model(X)
        loss = loss_fn(pred, y)

        # Backpropagation
        optimizer.zero_grad()
        loss.backward()
        
        # API call to trigger execution
        htcore.mark_step()
        
        optimizer.step()
        
        # API call to trigger execution
        htcore.mark_step()

        if batch % 100 == 0:
            loss, current = loss.item(), batch * len(X)
            print(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")

We also check the model’s performance against the test dataset to ensure it is learning.

def test(dataloader, model, loss_fn):
    size = len(dataloader.dataset)
def test(dataloader, model, loss_fn):
    size = len(dataloader.dataset)
    num_batches = len(dataloader)
    model.eval()
    test_loss, correct = 0, 0
    with torch.no_grad():
        for X, y in dataloader:
            X, y = X.to(device), y.to(device)
            pred = model(X)
            test_loss += loss_fn(pred, y).item()
            correct += (pred.argmax(1) == y).type(torch.float).sum().item()
    test_loss /= num_batches
    correct /= size
    print(f"Test Error: \n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \n")

The training process is conducted over several iterations (epochs). During each epoch, the model learns parameters to make better predictions. We print the model’s accuracy and loss at each epoch; we’d like to see the accuracy increase and the loss decrease with every epoch.

epochs = 5
for t in range(epochs):
    print(f"Epoch {t+1}\n-------------------------------")
    train(train_dataloader, model, loss_fn, optimizer)
    test(test_dataloader, model, loss_fn)
print("Done!")
Epoch 1 ------------------------------- loss: 2.310365 [ 0/60000] loss: 2.302217 [ 6400/60000] loss: 2.274406 [12800/60000] loss: 2.269496 [19200/60000] loss: 2.258406 [25600/60000] loss: 2.222645 [32000/60000] loss: 2.238458 [38400/60000] loss: 2.199006 [44800/60000] loss: 2.189651 [51200/60000] loss: 2.169230 [57600/60000] Test Error: Accuracy: 30.5%, Avg loss: 2.160159 Epoch 2 ------------------------------- loss: 2.167748 [ 0/60000] loss: 2.165102 [ 6400/60000] loss: 2.100914 [12800/60000] loss: 2.122778 [19200/60000] loss: 2.072059 [25600/60000] loss: 2.012739 [32000/60000] loss: 2.046291 [38400/60000] loss: 1.964602 [44800/60000] loss: 1.959184 [51200/60000] loss: 1.902106 [57600/60000] Test Error: Accuracy: 50.2%, Avg loss: 1.892953 Epoch 3 ------------------------------- loss: 1.922650 [ 0/60000] loss: 1.898303 [ 6400/60000] loss: 1.775073 [12800/60000] loss: 1.823592 [19200/60000] loss: 1.702000 [25600/60000] loss: 1.662501 [32000/60000] loss: 1.684734 [38400/60000] loss: 1.585656 [44800/60000] loss: 1.596274 [51200/60000] loss: 1.506514 [57600/60000] Test Error: Accuracy: 58.6%, Avg loss: 1.517501 Epoch 4 ------------------------------- loss: 1.579399 [ 0/60000] loss: 1.549755 [ 6400/60000] loss: 1.396002 [12800/60000] loss: 1.478985 [19200/60000] loss: 1.349092 [25600/60000] loss: 1.350576 [32000/60000] loss: 1.363445 [38400/60000] loss: 1.291087 [44800/60000] loss: 1.312230 [51200/60000] loss: 1.224336 [57600/60000] Test Error: Accuracy: 63.0%, Avg loss: 1.249705 Epoch 5 ------------------------------- loss: 1.321081 [ 0/60000] loss: 1.306476 [ 6400/60000] loss: 1.140665 [12800/60000] loss: 1.253876 [19200/60000] loss: 1.122847 [25600/60000] loss: 1.150012 [32000/60000] loss: 1.167040 [38400/60000] loss: 1.110087 [44800/60000] loss: 1.137658 [51200/60000] loss: 1.060612 [57600/60000] Test Error: Accuracy: 64.4%, Avg loss: 1.083993
Code language: CSS (css)

Done!

Read more about Training your model.

Saving Models

A common way to save a model is to serialize the internal state dictionary (containing the model parameters).

model = model.to("cpu")
torch.save(model.state_dict(), "model.pth")
print("Saved PyTorch Model State to model.pth")
Saved PyTorch Model State to model.pth
Code language: CSS (css)

Loading Models

The process for loading a model includes re-creating the model structure and loading the state dictionary into it.

model = NeuralNetwork()
model.load_state_dict(torch.load("model.pth"))

This model can now be used to make predictions.

classes = [
    "T-shirt/top",
    "Trouser",
    "Pullover",
    "Dress",
    "Coat",
    "Sandal",
    "Shirt",
    "Sneaker",
    "Bag",
    "Ankle boot",
]

model.eval()
x, y = test_data[0][0], test_data[0][1]
with torch.no_grad():
    pred = model(x)
    predicted, actual = classes[pred[0].argmax(0)], classes[y]
    print(f'Predicted: "{predicted}", Actual: "{actual}"')
Predicted: "Ankle boot", Actual: "Ankle boot"
Code language: JavaScript (javascript)

BSD 3-Clause License

Copyright (c) 2021 Habana Labs, Ltd. an Intel Company.
Copyright (c) 2017, Pytorch contributors
All rights reserved.

Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:
* Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
* Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution.
* Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission.
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS “AS IS” AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

Share this article:

Sign up for the latest Habana developer news, events, training, and updates.