数据处理
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 v2
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=v2.Compose([v2.ToImage(), v2.ToDtype(torch.float32, scale=True)]), ) # Download test data from open datasets. test_data = datasets.FashionMNIST( root="data", train=False, download=True, transform=v2.Compose([v2.ToImage(), v2.ToDtype(torch.float32, scale=True)]), )
0%| | 0.00/26.4M [00:00<?, ?B/s] 0%| | 65.5k/26.4M [00:00<01:09, 377kB/s] 1%| | 229k/26.4M [00:00<00:37, 707kB/s] 3%|▎ | 918k/26.4M [00:00<00:11, 2.18MB/s] 14%|█▍ | 3.67M/26.4M [00:00<00:03, 7.54MB/s] 37%|███▋ | 9.83M/26.4M [00:00<00:00, 17.5MB/s] 60%|██████ | 16.0M/26.4M [00:01<00:00, 23.5MB/s] 84%|████████▍ | 22.2M/26.4M [00:01<00:00, 27.4MB/s] 100%|██████████| 26.4M/26.4M [00:01<00:00, 20.1MB/s] 0%| | 0.00/29.5k [00:00<?, ?B/s] 100%|██████████| 29.5k/29.5k [00:00<00:00, 337kB/s] 0%| | 0.00/4.42M [00:00<?, ?B/s] 1%|▏ | 65.5k/4.42M [00:00<00:11, 375kB/s] 5%|▌ | 229k/4.42M [00:00<00:05, 708kB/s] 19%|█▉ | 852k/4.42M [00:00<00:01, 2.02MB/s] 79%|███████▊ | 3.47M/4.42M [00:00<00:00, 7.16MB/s] 100%|██████████| 4.42M/4.42M [00:00<00:00, 6.31MB/s] 0%| | 0.00/5.15k [00:00<?, ?B/s] 100%|██████████| 5.15k/5.15k [00:00<00:00, 50.8MB/s]
将Dataset传入DataLoader，可封装数据集的迭代功能，同时支持自动批处理、采样、打乱和多进程数据加载。这里定义批次大小为64，即数据加载器的每个迭代元素会返回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}")
breakShape of X [N, C, H, W]: torch.Size([64, 1, 28, 28]) Shape of y: torch.Size([64]) torch.int64
模型创建
在PyTorch中定义神经网络，需要创建一个继承自nn.Module的类。在__init__函数中定义网络层，在forward函数中指定数据在网络中的传播方式。为了加速神经网络运算，将模型迁移到CUDA、MPS、MTIA或XPU等加速器上；若当前加速器不可用，则使用CPU。
device = torch.accelerator.current_accelerator().type if torch.accelerator.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) ) )
模型参数优化
训练模型需要使用损失函数和优化器。
loss_fn = nn.CrossEntropyLoss() optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)
在单次训练循环中，模型会对训练数据集进行批量预测，并将预测误差反向传播以调整模型参数。
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}]")同时在测试数据集上验证模型性能，确保模型正在学习。
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")训练过程会进行多个迭代周期（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.308571 [ 64/60000] loss: 2.283845 [ 6464/60000] loss: 2.263912 [12864/60000] loss: 2.250840 [19264/60000] loss: 2.234350 [25664/60000] loss: 2.206552 [32064/60000] loss: 2.214474 [38464/60000] loss: 2.177198 [44864/60000] loss: 2.173431 [51264/60000] loss: 2.145613 [57664/60000] Test Error: Accuracy: 47.8%, Avg loss: 2.132438 Epoch 2 ------------------------------- loss: 2.151133 [ 64/60000] loss: 2.128767 [ 6464/60000] loss: 2.065868 [12864/60000] loss: 2.079103 [19264/60000] loss: 2.024301 [25664/60000] loss: 1.967980 [32064/60000] loss: 2.002589 [38464/60000] loss: 1.914585 [44864/60000] loss: 1.914785 [51264/60000] loss: 1.857888 [57664/60000] Test Error: Accuracy: 49.7%, Avg loss: 1.842189 Epoch 3 ------------------------------- loss: 1.880040 [ 64/60000] loss: 1.840284 [ 6464/60000] loss: 1.717671 [12864/60000] loss: 1.765090 [19264/60000] loss: 1.655026 [25664/60000] loss: 1.612522 [32064/60000] loss: 1.649840 [38464/60000] loss: 1.546005 [44864/60000] loss: 1.566916 [51264/60000] loss: 1.481809 [57664/60000] Test Error: Accuracy: 59.7%, Avg loss: 1.489499 Epoch 4 ------------------------------- loss: 1.551978 [ 64/60000] loss: 1.516600 [ 6464/60000] loss: 1.368065 [12864/60000] loss: 1.448740 [19264/60000] loss: 1.329922 [25664/60000] loss: 1.327002 [32064/60000] loss: 1.354688 [38464/60000] loss: 1.277455 [44864/60000] loss: 1.302673 [51264/60000] loss: 1.217838 [57664/60000] Test Error: Accuracy: 62.9%, Avg loss: 1.241062 Epoch 5 ------------------------------- loss: 1.309439 [ 64/60000] loss: 1.290670 [ 6464/60000] loss: 1.130989 [12864/60000] loss: 1.239824 [19264/60000] loss: 1.114016 [25664/60000] loss: 1.140699 [32064/60000] loss: 1.169178 [38464/60000] loss: 1.107137 [44864/60000] loss: 1.135468 [51264/60000] loss: 1.059944 [57664/60000] Test Error: Accuracy: 64.4%, Avg loss: 1.082769 Done!
模型保存
保存模型的常用方法是序列化包含模型参数的状态字典。
torch.save(model.state_dict(), "model.pth")
print("Saved PyTorch Model State to model.pth")Saved PyTorch Model State to model.pth
模型加载
加载模型的流程包括重新创建模型结构，并将状态字典加载到模型中。
model = NeuralNetwork().to(device)
model.load_state_dict(torch.load("model.pth", weights_only=True))<All keys matched successfully>
该模型可直接用于预测任务。
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"
