CoCalc Logo Icon
StoreFeaturesDocsShareSupportNewsAboutSign UpSign In
pytorch

Real-time collaboration for Jupyter Notebooks, Linux Terminals, LaTeX, VS Code, R IDE, and more,
all in one place. Commercial Alternative to JupyterHub.

GitHub Repository: pytorch/tutorials
Path: blob/main/beginner_source/basics/quickstart_tutorial.py
Views: 1017
1
"""
2
`Learn the Basics <intro.html>`_ ||
3
**Quickstart** ||
4
`Tensors <tensorqs_tutorial.html>`_ ||
5
`Datasets & DataLoaders <data_tutorial.html>`_ ||
6
`Transforms <transforms_tutorial.html>`_ ||
7
`Build Model <buildmodel_tutorial.html>`_ ||
8
`Autograd <autogradqs_tutorial.html>`_ ||
9
`Optimization <optimization_tutorial.html>`_ ||
10
`Save & Load Model <saveloadrun_tutorial.html>`_
11
12
Quickstart
13
===================
14
This section runs through the API for common tasks in machine learning. Refer to the links in each section to dive deeper.
15
16
Working with data
17
-----------------
18
PyTorch has two `primitives to work with data <https://pytorch.org/docs/stable/data.html>`_:
19
``torch.utils.data.DataLoader`` and ``torch.utils.data.Dataset``.
20
``Dataset`` stores the samples and their corresponding labels, and ``DataLoader`` wraps an iterable around
21
the ``Dataset``.
22
23
"""
24
25
import torch
26
from torch import nn
27
from torch.utils.data import DataLoader
28
from torchvision import datasets
29
from torchvision.transforms import ToTensor
30
31
######################################################################
32
# PyTorch offers domain-specific libraries such as `TorchText <https://pytorch.org/text/stable/index.html>`_,
33
# `TorchVision <https://pytorch.org/vision/stable/index.html>`_, and `TorchAudio <https://pytorch.org/audio/stable/index.html>`_,
34
# all of which include datasets. For this tutorial, we will be using a TorchVision dataset.
35
#
36
# The ``torchvision.datasets`` module contains ``Dataset`` objects for many real-world vision data like
37
# CIFAR, COCO (`full list here <https://pytorch.org/vision/stable/datasets.html>`_). In this tutorial, we
38
# use the FashionMNIST dataset. Every TorchVision ``Dataset`` includes two arguments: ``transform`` and
39
# ``target_transform`` to modify the samples and labels respectively.
40
41
# Download training data from open datasets.
42
training_data = datasets.FashionMNIST(
43
root="data",
44
train=True,
45
download=True,
46
transform=ToTensor(),
47
)
48
49
# Download test data from open datasets.
50
test_data = datasets.FashionMNIST(
51
root="data",
52
train=False,
53
download=True,
54
transform=ToTensor(),
55
)
56
57
######################################################################
58
# We pass the ``Dataset`` as an argument to ``DataLoader``. This wraps an iterable over our dataset, and supports
59
# automatic batching, sampling, shuffling and multiprocess data loading. Here we define a batch size of 64, i.e. each element
60
# in the dataloader iterable will return a batch of 64 features and labels.
61
62
batch_size = 64
63
64
# Create data loaders.
65
train_dataloader = DataLoader(training_data, batch_size=batch_size)
66
test_dataloader = DataLoader(test_data, batch_size=batch_size)
67
68
for X, y in test_dataloader:
69
print(f"Shape of X [N, C, H, W]: {X.shape}")
70
print(f"Shape of y: {y.shape} {y.dtype}")
71
break
72
73
######################################################################
74
# Read more about `loading data in PyTorch <data_tutorial.html>`_.
75
#
76
77
######################################################################
78
# --------------
79
#
80
81
################################
82
# Creating Models
83
# ------------------
84
# To define a neural network in PyTorch, we create a class that inherits
85
# from `nn.Module <https://pytorch.org/docs/stable/generated/torch.nn.Module.html>`_. We define the layers of the network
86
# in the ``__init__`` function and specify how data will pass through the network in the ``forward`` function. To accelerate
87
# operations in the neural network, we move it to the `accelerator <https://pytorch.org/docs/stable/torch.html#accelerators>`__
88
# such as CUDA, MPS, MTIA, or XPU. If the current accelerator is available, we will use it. Otherwise, we use the CPU.
89
90
device = torch.accelerator.current_accelerator().type if torch.accelerator.is_available() else "cpu"
91
print(f"Using {device} device")
92
93
# Define model
94
class NeuralNetwork(nn.Module):
95
def __init__(self):
96
super().__init__()
97
self.flatten = nn.Flatten()
98
self.linear_relu_stack = nn.Sequential(
99
nn.Linear(28*28, 512),
100
nn.ReLU(),
101
nn.Linear(512, 512),
102
nn.ReLU(),
103
nn.Linear(512, 10)
104
)
105
106
def forward(self, x):
107
x = self.flatten(x)
108
logits = self.linear_relu_stack(x)
109
return logits
110
111
model = NeuralNetwork().to(device)
112
print(model)
113
114
######################################################################
115
# Read more about `building neural networks in PyTorch <buildmodel_tutorial.html>`_.
116
#
117
118
119
######################################################################
120
# --------------
121
#
122
123
124
#####################################################################
125
# Optimizing the Model Parameters
126
# ----------------------------------------
127
# To train a model, we need a `loss function <https://pytorch.org/docs/stable/nn.html#loss-functions>`_
128
# and an `optimizer <https://pytorch.org/docs/stable/optim.html>`_.
129
130
loss_fn = nn.CrossEntropyLoss()
131
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)
132
133
134
#######################################################################
135
# In a single training loop, the model makes predictions on the training dataset (fed to it in batches), and
136
# backpropagates the prediction error to adjust the model's parameters.
137
138
def train(dataloader, model, loss_fn, optimizer):
139
size = len(dataloader.dataset)
140
model.train()
141
for batch, (X, y) in enumerate(dataloader):
142
X, y = X.to(device), y.to(device)
143
144
# Compute prediction error
145
pred = model(X)
146
loss = loss_fn(pred, y)
147
148
# Backpropagation
149
loss.backward()
150
optimizer.step()
151
optimizer.zero_grad()
152
153
if batch % 100 == 0:
154
loss, current = loss.item(), (batch + 1) * len(X)
155
print(f"loss: {loss:>7f} [{current:>5d}/{size:>5d}]")
156
157
##############################################################################
158
# We also check the model's performance against the test dataset to ensure it is learning.
159
160
def test(dataloader, model, loss_fn):
161
size = len(dataloader.dataset)
162
num_batches = len(dataloader)
163
model.eval()
164
test_loss, correct = 0, 0
165
with torch.no_grad():
166
for X, y in dataloader:
167
X, y = X.to(device), y.to(device)
168
pred = model(X)
169
test_loss += loss_fn(pred, y).item()
170
correct += (pred.argmax(1) == y).type(torch.float).sum().item()
171
test_loss /= num_batches
172
correct /= size
173
print(f"Test Error: \n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \n")
174
175
##############################################################################
176
# The training process is conducted over several iterations (*epochs*). During each epoch, the model learns
177
# parameters to make better predictions. We print the model's accuracy and loss at each epoch; we'd like to see the
178
# accuracy increase and the loss decrease with every epoch.
179
180
epochs = 5
181
for t in range(epochs):
182
print(f"Epoch {t+1}\n-------------------------------")
183
train(train_dataloader, model, loss_fn, optimizer)
184
test(test_dataloader, model, loss_fn)
185
print("Done!")
186
187
######################################################################
188
# Read more about `Training your model <optimization_tutorial.html>`_.
189
#
190
191
######################################################################
192
# --------------
193
#
194
195
######################################################################
196
# Saving Models
197
# -------------
198
# A common way to save a model is to serialize the internal state dictionary (containing the model parameters).
199
200
torch.save(model.state_dict(), "model.pth")
201
print("Saved PyTorch Model State to model.pth")
202
203
204
205
######################################################################
206
# Loading Models
207
# ----------------------------
208
#
209
# The process for loading a model includes re-creating the model structure and loading
210
# the state dictionary into it.
211
212
model = NeuralNetwork().to(device)
213
model.load_state_dict(torch.load("model.pth", weights_only=True))
214
215
#############################################################
216
# This model can now be used to make predictions.
217
218
classes = [
219
"T-shirt/top",
220
"Trouser",
221
"Pullover",
222
"Dress",
223
"Coat",
224
"Sandal",
225
"Shirt",
226
"Sneaker",
227
"Bag",
228
"Ankle boot",
229
]
230
231
model.eval()
232
x, y = test_data[0][0], test_data[0][1]
233
with torch.no_grad():
234
x = x.to(device)
235
pred = model(x)
236
predicted, actual = classes[pred[0].argmax(0)], classes[y]
237
print(f'Predicted: "{predicted}", Actual: "{actual}"')
238
239
240
######################################################################
241
# Read more about `Saving & Loading your model <saveloadrun_tutorial.html>`_.
242
#
243
244