PyTorch Quickstart

欢迎来到啾啾的博客🐱。
记录学习点滴。分享工作思考和实用技巧，偶尔也分享一些杂谈💬。
有很多很多不足的地方，欢迎评论交流，感谢您的阅读和评论😄。

@[TOC

引言

本篇较官网的更清晰，带你快速入门Quickstart。

资料：
PyTorch Quick Star

1 数据处理

1.1 DataLoader与Dataset

首先需要了解 torch.utils.data.DataLoader 和 torch.utils.data.Dataset。

Dataset像“仓库管理员” ,只负责两件事：

• 告诉 PyTorch 一共有多少条样本（__len__）。
• 按索引把一条样本和它的标签取出来（__getitem__）。
  不管数据是图片、文本还是表格，只要继承 torch.utils.data.Dataset 并实现这两个方法，就能被 PyTorch 识别。

class MyDataset(torch.utils.data.Dataset):
    def __init__(self, csv_file):
        self.data = pd.read_csv(csv_file)
    def __len__(self):
        return len(self.data)
    def __getitem__(self, idx):
        x = self.data.iloc[idx, :-1].values
        y = self.data.iloc[idx, -1]
        return torch.tensor(x, dtype=torch.float32), torch.tensor(y, dtype=torch.long)

DataLoader：像“流水线操作员” 。把上面那个 Dataset 包起来，帮你“批量、乱序、并行”地把样本送进模型。 常用参数：

• batch_size：一次抓多少条样本。
• shuffle：每个 epoch 是否打乱顺序。
• num_workers：用多少个子进程并行加载数据，避免 GPU 等 CPU。

loader = torch.utils.data.DataLoader(MyDataset('train.csv'),
                                     batch_size=64,
                                     shuffle=True,
                                     num_workers=4)
for x, y in loader:        # 每次循环拿到一个 64×… 的 batch
    pred = model(x)
    loss = crit(pred, y)
    ...

1.2 领域工具箱

同时，PyTorch 提供了特定领域的库，如TorchText、 TorchVision和 TorchAudio，这些库相当于PyTorch 官方出的三个“领域工具箱”，每个工具箱都帮你把该领域最常用的数据、预处理、模型和指标一次性打包好了；你不再需要自己从零去爬数据、写 transform、找权重。

就是官方帮你把 NLP、CV、语音三大领域“数据集 + 预处理 + 经典模型 + 预训练权重”全部打包好的扩展库，让你用最少的代码把“数据→模型→训练”跑通。

1. TorchText：自然语言处理（NLP）工具箱
   • 内置数据集：IMDb、SST、WikiText2、Multi30k 等。
   • 文本预处理：分词器（tokenizer）、SentencePiece、BPE、词表构建、截断/填充、数值化。
   • 预训练向量：GloVe、FastText、CharNGram。
   • 经典模型/组件：RoBERTa、XLM-R 等封装好的 encoder bundle。

2. TorchVision：计算机视觉（CV）工具箱
   • 内置数据集：CIFAR-10/100、ImageNet、COCO、Cityscapes、MNIST、CelebA …
   • 图像/视频增广：旋转、裁剪、颜色抖动、MixUp、CutMix、AutoAugment。
   • 预训练模型：ResNet、EfficientNet、YOLO、Mask R-CNN、ViT、Swin Transformer 等。
   • 工具函数：画框、画分割掩膜、可视化、视频编解码。

3. TorchAudio：语音与音频处理工具箱
   • 内置数据集：LibriSpeech、CommonVoice、VoxCeleb、YESNO、GTZAN 等。
   • 音频 I/O：读写 wav、mp3、flac；支持 FFmpeg 后端。
   • 信号处理：STFT、Mel 频谱、MFCC、滤波、重采样、变速/变调。
   • 预训练模型/流水线：Wav2Vec 2.0、HuBERT、Whisper、CTC 解码器、ASR/增强/分离 pipeline。

下载数据并封装成DataLoader

import torch  
from torch import nn  
from torch.utils.data import DataLoader  
from torchvision import datasets  
from torchvision.transforms import ToTensor  
  
# 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(),  
)  
  
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

2 优化模型参数

在之前通过基础类构建的MLP中，进行“更新参数”过程需要使用超参数手动更新。

loss.backward()

with torch.no_grad():
    for param in model.parameters():
        param -= 1e-3 * param.grad  # 手动更新：学习率 × 梯度
    model.zero_grad()

一个损失函数、一个优化器也可以。

loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)

# 步骤
loss.backward()        # 计算梯度
optimizer.step()       # ← 自动完成“用梯度更新参数”
optimizer.zero_grad()  # 清空梯度

• 常见优化器
  • SGD：简单按梯度方向走一步
  • Adam：自适应学习率，对每个参数调整步长
  • RMSprop：缓解梯度震荡

在一个训练循环中，模型对训练数据集（以批处理方式输入）进行预测，并将预测误差反向传播以调整模型的参数。

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

训练过程在多个迭代（周期）中进行。在每个周期中，模型学习参数以做出更好的预测。我们在每个周期打印模型的准确率和损失；我们希望看到准确率随着每个周期的进行而提高，损失则逐渐降低。

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!")

3 保存模型与加载使用

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

加载模型的流程包括重新创建模型结构并将其状态字典加载进去。

# 保存模型
torch.save(model.state_dict(), "model.pth")
print("Saved PyTorch Model State to model.pth")

# 重新加载模型
model = NeuralNetwork().to(device)
model.load_state_dict(torch.load("model.pth", weights_only=True))

# 使用模型进行预测
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}"')

4 完整Quicksatrt Demo

import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor

# 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(),
)

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


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)

# 损失函数和一个优化器
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")

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!")


# 保存模型
torch.save(model.state_dict(), "model.pth")
print("Saved PyTorch Model State to model.pth")

# 重新加载模型
model = NeuralNetwork().to(device)
model.load_state_dict(torch.load("model.pth", weights_only=True))

# 使用模型进行预测
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}"')