Quickstart — PyTorch Tutorials 2.3.0+cu121 documentation (2024)

Tutorials >
Quickstart

Shortcuts

beginner/basics/quickstart_tutorial

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

Note

Working with data¶

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

import torchfrom torch import nnfrom torch.utils.data import DataLoaderfrom torchvision import datasetsfrom torchvision.transforms import ToTensor

PyTorch offers domain-specific libraries such as TorchText,TorchVision, 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 likeCIFAR, COCO (full list here). In this tutorial, weuse the FashionMNIST dataset. Every TorchVision Dataset includes two arguments: transform andtarget_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(),)

Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/train-images-idx3-ubyte.gzDownloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/train-images-idx3-ubyte.gz to data/FashionMNIST/raw/train-images-idx3-ubyte.gz 0%| | 0/26421880 [00:00<?, ?it/s] 0%| | 65536/26421880 [00:00<01:12, 364115.76it/s] 1%| | 229376/26421880 [00:00<00:38, 681379.99it/s] 4%|3 | 950272/26421880 [00:00<00:11, 2184680.52it/s] 15%|#4 | 3833856/26421880 [00:00<00:02, 7587134.53it/s] 31%|###1 | 8257536/26421880 [00:00<00:01, 16267860.73it/s] 45%|####4 | 11862016/26421880 [00:00<00:00, 17951766.93it/s] 63%|######3 | 16646144/26421880 [00:01<00:00, 24933184.46it/s] 78%|#######7 | 20578304/26421880 [00:01<00:00, 24103564.46it/s] 97%|#########7| 25690112/26421880 [00:01<00:00, 30368939.72it/s]100%|##########| 26421880/26421880 [00:01<00:00, 18163618.83it/s]Extracting data/FashionMNIST/raw/train-images-idx3-ubyte.gz to data/FashionMNIST/rawDownloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/train-labels-idx1-ubyte.gzDownloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/train-labels-idx1-ubyte.gz to data/FashionMNIST/raw/train-labels-idx1-ubyte.gz 0%| | 0/29515 [00:00<?, ?it/s]100%|##########| 29515/29515 [00:00<00:00, 331054.58it/s]Extracting data/FashionMNIST/raw/train-labels-idx1-ubyte.gz to data/FashionMNIST/rawDownloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/t10k-images-idx3-ubyte.gzDownloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/t10k-images-idx3-ubyte.gz to data/FashionMNIST/raw/t10k-images-idx3-ubyte.gz 0%| | 0/4422102 [00:00<?, ?it/s] 1%|1 | 65536/4422102 [00:00<00:11, 365606.92it/s] 5%|5 | 229376/4422102 [00:00<00:06, 687107.27it/s] 21%|## | 917504/4422102 [00:00<00:01, 2123399.43it/s] 70%|######9 | 3080192/4422102 [00:00<00:00, 6653714.36it/s]100%|##########| 4422102/4422102 [00:00<00:00, 6122626.60it/s]Extracting data/FashionMNIST/raw/t10k-images-idx3-ubyte.gz to data/FashionMNIST/rawDownloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/t10k-labels-idx1-ubyte.gzDownloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/t10k-labels-idx1-ubyte.gz to data/FashionMNIST/raw/t10k-labels-idx1-ubyte.gz 0%| | 0/5148 [00:00<?, ?it/s]100%|##########| 5148/5148 [00:00<00:00, 39258685.44it/s]Extracting data/FashionMNIST/raw/t10k-labels-idx1-ubyte.gz to data/FashionMNIST/raw

We pass the Dataset as an argument to DataLoader. This wraps an iterable over our dataset, and supportsautomatic batching, sampling, shuffling and multiprocess data loading. Here we define a batch size of 64, i.e. each elementin 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(f"Shape of X [N, C, H, W]: {X.shape}") print(f"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

Read more about loading data in PyTorch.

Creating Models¶

To define a neural network in PyTorch, we create a class that inheritsfrom nn.Module. We define the layers of the networkin the __init__ function and specify how data will pass through the network in the forward function. To accelerateoperations in the neural network, we move it to the GPU or MPS if available.

# Get cpu, gpu or mps device for training.device = ( "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu")print(f"Using {device} device")# Define modelclass NeuralNetwork(nn.Module): def __init__(self): super().__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 logitsmodel = NeuralNetwork().to(device)print(model)

Optimizing the Model Parameters¶

To train a model, we need a loss functionand 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), andbackpropagates 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 loss.backward() optimizer.step() optimizer.zero_grad() if batch % 100 == 0: loss, current = loss.item(), (batch + 1) * 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) 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 learnsparameters to make better predictions. We print the model’s accuracy and loss at each epoch; we’d like to see theaccuracy increase and the loss decrease with every epoch.

epochs = 5for 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.303494 [ 64/60000]loss: 2.294637 [ 6464/60000]loss: 2.277102 [12864/60000]loss: 2.269977 [19264/60000]loss: 2.254235 [25664/60000]loss: 2.237146 [32064/60000]loss: 2.231055 [38464/60000]loss: 2.205037 [44864/60000]loss: 2.203240 [51264/60000]loss: 2.170889 [57664/60000]Test Error: Accuracy: 53.9%, Avg loss: 2.168588Epoch 2-------------------------------loss: 2.177787 [ 64/60000]loss: 2.168083 [ 6464/60000]loss: 2.114910 [12864/60000]loss: 2.130412 [19264/60000]loss: 2.087473 [25664/60000]loss: 2.039670 [32064/60000]loss: 2.054274 [38464/60000]loss: 1.985457 [44864/60000]loss: 1.996023 [51264/60000]loss: 1.917241 [57664/60000]Test Error: Accuracy: 60.2%, Avg loss: 1.920374Epoch 3-------------------------------loss: 1.951705 [ 64/60000]loss: 1.919516 [ 6464/60000]loss: 1.808730 [12864/60000]loss: 1.846550 [19264/60000]loss: 1.740618 [25664/60000]loss: 1.698733 [32064/60000]loss: 1.708889 [38464/60000]loss: 1.614436 [44864/60000]loss: 1.646475 [51264/60000]loss: 1.524308 [57664/60000]Test Error: Accuracy: 61.4%, Avg loss: 1.547092Epoch 4-------------------------------loss: 1.612695 [ 64/60000]loss: 1.570870 [ 6464/60000]loss: 1.424730 [12864/60000]loss: 1.489542 [19264/60000]loss: 1.367256 [25664/60000]loss: 1.373464 [32064/60000]loss: 1.376744 [38464/60000]loss: 1.304962 [44864/60000]loss: 1.347154 [51264/60000]loss: 1.230661 [57664/60000]Test Error: Accuracy: 62.7%, Avg loss: 1.260891Epoch 5-------------------------------loss: 1.337803 [ 64/60000]loss: 1.313278 [ 6464/60000]loss: 1.151837 [12864/60000]loss: 1.252142 [19264/60000]loss: 1.123048 [25664/60000]loss: 1.159531 [32064/60000]loss: 1.175011 [38464/60000]loss: 1.115554 [44864/60000]loss: 1.160974 [51264/60000]loss: 1.062730 [57664/60000]Test Error: Accuracy: 64.6%, Avg loss: 1.087374Done!

Saving Models¶

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

torch.save(model.state_dict(), "model.pth")print("Saved PyTorch Model State to model.pth")

Saved PyTorch Model State to model.pth

Loading Models¶

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

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

<All keys matched successfully>

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(): x = x.to(device) pred = model(x) predicted, actual = classes[pred[0].argmax(0)], classes[y] print(f'Predicted: "{predicted}", Actual: "{actual}"')

Predicted: "Ankle boot", Actual: "Ankle boot"