[PyTorch 2.0 翻译&学习] 0. 快速开始

Quickstart

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 data: torch.utils.data.DataLoader and torch.utils.data.Dataset. Dataset stores the samples and their corresponding labels, and DataLoader wraps an iterable around the Dataset.

快速开始

本节将介绍机器学习中常见任务的API。请参考每个部分中的链接以深入了解。

使用数据

PyTorch 有两个用于处理数据的基本类：torch.utils.data.DataLoader 和 torch.utils.data.Dataset。Dataset 存储样本及其对应的标签，DataLoader 包装了 Dataset，使其成为一个可迭代对象。

import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets
from 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 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.

PyTorch 提供了领域特定的库，如 TorchText、TorchVision 和 TorchAudio，这些库都包含了数据集。在本教程中，我们将使用 TorchVision 数据集。

torchvision.datasets 模块包含许多真实世界视觉数据的 Dataset 对象，如 CIFAR、COCO（完整列表请参见此处）。在本教程中，我们将使用 FashionMNIST 数据集。每个 TorchVision Dataset 都包含两个参数：transform 和 target_transform，用于分别修改样本和标签。

# 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.gz
Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/train-images-idx3-ubyte.gz to data/FashionMNIST/raw/train-images-idx3-ubyte.gz


100.0%


Extracting data/FashionMNIST/raw/train-images-idx3-ubyte.gz to data/FashionMNIST/raw

Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/train-labels-idx1-ubyte.gz
Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/train-labels-idx1-ubyte.gz to data/FashionMNIST/raw/train-labels-idx1-ubyte.gz


100.0%


Extracting data/FashionMNIST/raw/train-labels-idx1-ubyte.gz to data/FashionMNIST/raw

Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/t10k-images-idx3-ubyte.gz
Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/t10k-images-idx3-ubyte.gz to data/FashionMNIST/raw/t10k-images-idx3-ubyte.gz


100.0%


Extracting data/FashionMNIST/raw/t10k-images-idx3-ubyte.gz to data/FashionMNIST/raw

Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/t10k-labels-idx1-ubyte.gz
Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/t10k-labels-idx1-ubyte.gz to data/FashionMNIST/raw/t10k-labels-idx1-ubyte.gz


100.0%


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 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.

我们将 Dataset 作为参数传递给 DataLoader。DataLoader 包装了我们的数据集，支持自动批处理、采样、洗牌和多进程数据加载。在这里，我们定义了一个批处理大小为 64，即 DataLoader 可迭代对象中的每个元素将返回一个包含 64 个特征和标签的批次。

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.

读取更多关于在 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 the GPU or MPS if available.

创建模型

要在 PyTorch 中定义神经网络，我们创建一个从 nn.Module 继承的类。我们在 init 函数中定义网络的层，并在 forward 函数中指定数据如何通过网络。为了加速神经网络中的操作，我们将其移动到 GPU 或 MPS（如果可用）。

# 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 model
class 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 logits

model = NeuralNetwork().to(device)
print(model)

Using cuda 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)
  )
)

Read more about building neural networks in PyTorch.

读取更多关于在 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
        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 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.

训练过程在几个迭代（epoch）中进行。在每个 epoch 中，模型学习参数以做出更好的预测。我们在每个 epoch 打印模型的准确率和损失；我们希望看到准确率随着每个 epoch 的增加而增加，损失随着每个 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.298584  [   64/60000]
loss: 2.287438  [ 6464/60000]
loss: 2.272034  [12864/60000]
loss: 2.272073  [19264/60000]
loss: 2.248833  [25664/60000]
loss: 2.224293  [32064/60000]
loss: 2.223804  [38464/60000]
loss: 2.195088  [44864/60000]
loss: 2.199643  [51264/60000]
loss: 2.158751  [57664/60000]
Test Error: 
 Accuracy: 52.0%, Avg loss: 2.159428 

Epoch 2
-------------------------------
loss: 2.168097  [   64/60000]
loss: 2.159347  [ 6464/60000]
loss: 2.097428  [12864/60000]
loss: 2.121088  [19264/60000]
loss: 2.069698  [25664/60000]
loss: 2.006833  [32064/60000]
loss: 2.033619  [38464/60000]
loss: 1.952340  [44864/60000]
loss: 1.964535  [51264/60000]
loss: 1.877724  [57664/60000]
Test Error: 
 Accuracy: 54.6%, Avg loss: 1.888157 

Epoch 3
-------------------------------
loss: 1.919495  [   64/60000]
loss: 1.891491  [ 6464/60000]
loss: 1.765739  [12864/60000]
loss: 1.814736  [19264/60000]
loss: 1.709010  [25664/60000]
loss: 1.650090  [32064/60000]
loss: 1.673284  [38464/60000]
loss: 1.571386  [44864/60000]
loss: 1.602315  [51264/60000]
loss: 1.484661  [57664/60000]
Test Error: 
 Accuracy: 61.3%, Avg loss: 1.513954 

Epoch 4
-------------------------------
loss: 1.579137  [   64/60000]
loss: 1.547676  [ 6464/60000]
loss: 1.390591  [12864/60000]
loss: 1.472600  [19264/60000]
loss: 1.358534  [25664/60000]
loss: 1.346502  [32064/60000]
loss: 1.361573  [38464/60000]
loss: 1.282488  [44864/60000]
loss: 1.323595  [51264/60000]
loss: 1.216870  [57664/60000]
Test Error: 
 Accuracy: 63.8%, Avg loss: 1.248743 

Epoch 5
-------------------------------
loss: 1.323229  [   64/60000]
loss: 1.308202  [ 6464/60000]
loss: 1.135307  [12864/60000]
loss: 1.253654  [19264/60000]
loss: 1.134068  [25664/60000]
loss: 1.152532  [32064/60000]
loss: 1.173257  [38464/60000]
loss: 1.105085  [44864/60000]
loss: 1.152041  [51264/60000]
loss: 1.059873  [57664/60000]
Test Error: 
 Accuracy: 64.9%, Avg loss: 1.085592 

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).

保存模型

保存模型的常见方法是序列化内部状态字典（包含模型参数）。

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 loading the 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"

Read more about Saving & Loading your model.

读取更多关于保存和加载模型的信息。