Real-time collaboration for Jupyter Notebooks, Linux Terminals, LaTeX, VS Code, R IDE, and more,
all in one place. Commercial Alternative to JupyterHub.
Real-time collaboration for Jupyter Notebooks, Linux Terminals, LaTeX, VS Code, R IDE, and more,
all in one place. Commercial Alternative to JupyterHub.
Path: blob/main/beginner_source/basics/quickstart_tutorial.py
Views: 1017
"""1`Learn the Basics <intro.html>`_ ||2**Quickstart** ||3`Tensors <tensorqs_tutorial.html>`_ ||4`Datasets & DataLoaders <data_tutorial.html>`_ ||5`Transforms <transforms_tutorial.html>`_ ||6`Build Model <buildmodel_tutorial.html>`_ ||7`Autograd <autogradqs_tutorial.html>`_ ||8`Optimization <optimization_tutorial.html>`_ ||9`Save & Load Model <saveloadrun_tutorial.html>`_1011Quickstart12===================13This section runs through the API for common tasks in machine learning. Refer to the links in each section to dive deeper.1415Working with data16-----------------17PyTorch has two `primitives to work with data <https://pytorch.org/docs/stable/data.html>`_:18``torch.utils.data.DataLoader`` and ``torch.utils.data.Dataset``.19``Dataset`` stores the samples and their corresponding labels, and ``DataLoader`` wraps an iterable around20the ``Dataset``.2122"""2324import torch25from torch import nn26from torch.utils.data import DataLoader27from torchvision import datasets28from torchvision.transforms import ToTensor2930######################################################################31# PyTorch offers domain-specific libraries such as `TorchText <https://pytorch.org/text/stable/index.html>`_,32# `TorchVision <https://pytorch.org/vision/stable/index.html>`_, and `TorchAudio <https://pytorch.org/audio/stable/index.html>`_,33# all of which include datasets. For this tutorial, we will be using a TorchVision dataset.34#35# The ``torchvision.datasets`` module contains ``Dataset`` objects for many real-world vision data like36# CIFAR, COCO (`full list here <https://pytorch.org/vision/stable/datasets.html>`_). In this tutorial, we37# use the FashionMNIST dataset. Every TorchVision ``Dataset`` includes two arguments: ``transform`` and38# ``target_transform`` to modify the samples and labels respectively.3940# Download training data from open datasets.41training_data = datasets.FashionMNIST(42root="data",43train=True,44download=True,45transform=ToTensor(),46)4748# Download test data from open datasets.49test_data = datasets.FashionMNIST(50root="data",51train=False,52download=True,53transform=ToTensor(),54)5556######################################################################57# We pass the ``Dataset`` as an argument to ``DataLoader``. This wraps an iterable over our dataset, and supports58# automatic batching, sampling, shuffling and multiprocess data loading. Here we define a batch size of 64, i.e. each element59# in the dataloader iterable will return a batch of 64 features and labels.6061batch_size = 646263# Create data loaders.64train_dataloader = DataLoader(training_data, batch_size=batch_size)65test_dataloader = DataLoader(test_data, batch_size=batch_size)6667for X, y in test_dataloader:68print(f"Shape of X [N, C, H, W]: {X.shape}")69print(f"Shape of y: {y.shape} {y.dtype}")70break7172######################################################################73# Read more about `loading data in PyTorch <data_tutorial.html>`_.74#7576######################################################################77# --------------78#7980################################81# Creating Models82# ------------------83# To define a neural network in PyTorch, we create a class that inherits84# from `nn.Module <https://pytorch.org/docs/stable/generated/torch.nn.Module.html>`_. We define the layers of the network85# in the ``__init__`` function and specify how data will pass through the network in the ``forward`` function. To accelerate86# operations in the neural network, we move it to the `accelerator <https://pytorch.org/docs/stable/torch.html#accelerators>`__87# such as CUDA, MPS, MTIA, or XPU. If the current accelerator is available, we will use it. Otherwise, we use the CPU.8889device = torch.accelerator.current_accelerator().type if torch.accelerator.is_available() else "cpu"90print(f"Using {device} device")9192# Define model93class NeuralNetwork(nn.Module):94def __init__(self):95super().__init__()96self.flatten = nn.Flatten()97self.linear_relu_stack = nn.Sequential(98nn.Linear(28*28, 512),99nn.ReLU(),100nn.Linear(512, 512),101nn.ReLU(),102nn.Linear(512, 10)103)104105def forward(self, x):106x = self.flatten(x)107logits = self.linear_relu_stack(x)108return logits109110model = NeuralNetwork().to(device)111print(model)112113######################################################################114# Read more about `building neural networks in PyTorch <buildmodel_tutorial.html>`_.115#116117118######################################################################119# --------------120#121122123#####################################################################124# Optimizing the Model Parameters125# ----------------------------------------126# To train a model, we need a `loss function <https://pytorch.org/docs/stable/nn.html#loss-functions>`_127# and an `optimizer <https://pytorch.org/docs/stable/optim.html>`_.128129loss_fn = nn.CrossEntropyLoss()130optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)131132133#######################################################################134# In a single training loop, the model makes predictions on the training dataset (fed to it in batches), and135# backpropagates the prediction error to adjust the model's parameters.136137def train(dataloader, model, loss_fn, optimizer):138size = len(dataloader.dataset)139model.train()140for batch, (X, y) in enumerate(dataloader):141X, y = X.to(device), y.to(device)142143# Compute prediction error144pred = model(X)145loss = loss_fn(pred, y)146147# Backpropagation148loss.backward()149optimizer.step()150optimizer.zero_grad()151152if batch % 100 == 0:153loss, current = loss.item(), (batch + 1) * len(X)154print(f"loss: {loss:>7f} [{current:>5d}/{size:>5d}]")155156##############################################################################157# We also check the model's performance against the test dataset to ensure it is learning.158159def test(dataloader, model, loss_fn):160size = len(dataloader.dataset)161num_batches = len(dataloader)162model.eval()163test_loss, correct = 0, 0164with torch.no_grad():165for X, y in dataloader:166X, y = X.to(device), y.to(device)167pred = model(X)168test_loss += loss_fn(pred, y).item()169correct += (pred.argmax(1) == y).type(torch.float).sum().item()170test_loss /= num_batches171correct /= size172print(f"Test Error: \n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \n")173174##############################################################################175# The training process is conducted over several iterations (*epochs*). During each epoch, the model learns176# parameters to make better predictions. We print the model's accuracy and loss at each epoch; we'd like to see the177# accuracy increase and the loss decrease with every epoch.178179epochs = 5180for t in range(epochs):181print(f"Epoch {t+1}\n-------------------------------")182train(train_dataloader, model, loss_fn, optimizer)183test(test_dataloader, model, loss_fn)184print("Done!")185186######################################################################187# Read more about `Training your model <optimization_tutorial.html>`_.188#189190######################################################################191# --------------192#193194######################################################################195# Saving Models196# -------------197# A common way to save a model is to serialize the internal state dictionary (containing the model parameters).198199torch.save(model.state_dict(), "model.pth")200print("Saved PyTorch Model State to model.pth")201202203204######################################################################205# Loading Models206# ----------------------------207#208# The process for loading a model includes re-creating the model structure and loading209# the state dictionary into it.210211model = NeuralNetwork().to(device)212model.load_state_dict(torch.load("model.pth", weights_only=True))213214#############################################################215# This model can now be used to make predictions.216217classes = [218"T-shirt/top",219"Trouser",220"Pullover",221"Dress",222"Coat",223"Sandal",224"Shirt",225"Sneaker",226"Bag",227"Ankle boot",228]229230model.eval()231x, y = test_data[0][0], test_data[0][1]232with torch.no_grad():233x = x.to(device)234pred = model(x)235predicted, actual = classes[pred[0].argmax(0)], classes[y]236print(f'Predicted: "{predicted}", Actual: "{actual}"')237238239######################################################################240# Read more about `Saving & Loading your model <saveloadrun_tutorial.html>`_.241#242243244