PyTorch：像写 Python 一样写深度学习

# python first_grad.py
import torch
import torch.nn as nn

model = nn.Linear(4, 2)              # 一层全连接：4 维进，2 维出
x = torch.randn(3, 4)                # 3 条 4 维样本
y = model(x)                         # 前向传播
print(y.shape)                       # torch.Size([3, 2])

loss = (y ** 2).mean()               # 随手定义一个标量损失
loss.backward()                      # 反向传播：梯度已写进参数的 .grad
print(model.weight.grad.shape)       # torch.Size([2, 4])

>>> import torch, numpy as np
>>> t = torch.tensor([[1.0, 2.0], [3.0, 4.0]])    # 从数据建
>>> r = torch.randn(2, 3)                          # 正态随机，参数初始化常用
>>> t.shape, t.dtype, t.device
(torch.Size([2, 2]), torch.float32, device(type='cpu'))

>>> a = np.ones((2, 2))
>>> n = torch.from_numpy(a)                        # 与 numpy 共享内存！
>>> n[0, 0] = 9.0; a[0, 0]
9.0
>>> n.numpy()                                      # 转回 numpy（GPU 上要先 t.cpu()）

>>> x = torch.tensor(3.0, requires_grad=True)
>>> y = x ** 2 + 2 * x                    # y = x² + 2x，导数 2x + 2
>>> y.backward()
>>> x.grad                                 # x=3 处应为 8
tensor(8.)
>>> x.grad.zero_()                        # 下一次 backward 前清零——
# .grad 是累加的，不清零梯度会越滚越大，这是新手第一大坑（见第 07 节）

import torch.nn as nn

class MLP(nn.Module):
    def __init__(self):
        super().__init__()               # 必调：初始化注册表
        self.net = nn.Sequential(        # 按顺序组一层串
            nn.Flatten(),                # (N, 1, 28, 28) -> (N, 784)
            nn.Linear(784, 256),
            nn.ReLU(),
            nn.Linear(256, 10),
        )

    def forward(self, x):                # 定义前向；model(x) 即调用它
        return self.net(x)

model = MLP()
print(sum(p.numel() for p in model.parameters()))   # 203530（权重 + 偏置都算上）

from torch.utils.data import Dataset, DataLoader

class MyData(Dataset):
    def __len__(self): ...        # 样本总数
    def __getitem__(self, i): ... # 第 i 条样本，返回 (特征, 标签)

loader = DataLoader(MyData(), batch_size=64, shuffle=True)

for epoch in range(num_epochs):
    for x, y in loader:
        pred = model(x)               # 1 前向
        optimizer.zero_grad()         # 2 清零（必须在 backward 前）
        loss = criterion(pred, y)
        loss.backward()               # 3 反向
        optimizer.step()              # 4 更新
        running += loss.item()        # 5 记录（.item() 取出 python 数字）

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

transform = transforms.ToTensor()                 # 像素 0~255 转 0~1 张量
train_ds = datasets.MNIST("data", train=True,  download=True, transform=transform)
test_ds  = datasets.MNIST("data", train=False, download=True, transform=transform)
train_loader = DataLoader(train_ds, batch_size=64, shuffle=True)
test_loader  = DataLoader(test_ds,  batch_size=256)

# 设备无关写法：有 CUDA 用 CUDA，否则 CPU（Apple Silicon 可改成 "mps"）
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Flatten(),                          # (N, 1, 28, 28) -> (N, 784)
            nn.Linear(784, 128),
            nn.ReLU(),
            nn.Dropout(0.2),                       # 防过拟合
            nn.Linear(128, 10),                    # 输出 10 类的 logits
        )

    def forward(self, x):
        return self.net(x)

model = MLP().to(device)                           # 模型搬上设备
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

for epoch in range(3):
    model.train()                                  # 训练模式：Dropout 生效
    running = 0.0
    for x, y in train_loader:
        x, y = x.to(device), y.to(device)          # 数据搬上同一设备

        pred = model(x)                            # 1 前向
        optimizer.zero_grad()                      # 2 清零（不清会累积）
        loss = criterion(pred, y)
        loss.backward()                            # 3 反向
        optimizer.step()                           # 4 更新

        running += loss.item() * len(y)            # 5 记录本批 loss

    print(f"epoch {epoch + 1}  loss {running / len(train_ds):.4f}")

# epoch 1  loss 0.3901（示意输出，数值随随机初始化波动）
# epoch 2  loss 0.1967（示意输出）
# epoch 3  loss 0.1453（示意输出）

model.eval()                                       # 评估模式：Dropout 关闭
correct = 0
with torch.no_grad():                              # 不建计算图：省内存、更快
    for x, y in test_loader:
        x, y = x.to(device), y.to(device)
        pred = model(x).argmax(dim=1)              # logits 最大者当预测类别
        correct += (pred == y).sum().item()

print(f"测试集准确率 {correct / len(test_ds):.3f}")    # 0.97x 左右（示意输出）

# 保存：只存参数字典（官方推荐），换环境加载不依赖代码路径
torch.save(model.state_dict(), "mnist_mlp.pth")

# 加载：先有同结构模型，再灌参数
model2 = MLP().to(device)
model2.load_state_dict(torch.load("mnist_mlp.pth", map_location=device))
model2.eval()

# CUDA / MPS / CPU 三选一，替换第 1 步里的 device 那行即可
if torch.cuda.is_available():
    device = torch.device("cuda")
elif torch.backends.mps.is_available():
    device = torch.device("mps")
else:
    device = torch.device("cpu")

# model.to(device)、x.to(device) 原样保留——这就是 device 无关写法的意义

# 错误写法：只有 forward + backward，梯度永远在叠加
pred = model(x); loss.backward(); optimizer.step()

# 正确：backward 之前清零（set_to_none=True 是默认行为，直接释放旧梯度）
optimizer.zero_grad()          # 放 backward 之前，一步都不能省
pred = model(x); loss.backward(); optimizer.step()

model.eval()
with torch.no_grad():
    ...                      # 评估、推理都放这里
model.train()                # 回到训练循环前切回来，别一路 eval 到底

model = MLP().to(device)                    # 模型上设备
x, y = x.to(device), y.to(device)           # 输入和标签也上同一设备

result = model(x)
result.cpu().numpy()                        # GPU 张量先 .cpu() 再转 numpy

# 排查口诀：print(t.device) 逐个看，找到 cpu 的那个搬过去

# 不推荐：torch.save(model, "m.pt") —— 绑死类的定义位置

# 推荐：存 state_dict，加载方自己 import 模型类
torch.save(model.state_dict(), "mnist_mlp.pth")
model2 = MLP()                                          # 结构在本项目里现建
model2.load_state_dict(torch.load("mnist_mlp.pth", map_location=device))
model2.eval()

# 推理场景没有梯度，load 时配 map_location 指定设备，跨机器最稳