2、PyTorch框架入门：张量操作与自动求导、构建简单全连接网络、数据加载器与训练循环、模型保存与加载
说实话，很多初学者一上来就扎进Transformer、大模型，结果连梯度怎么回传都搞不清楚。我个人觉得，PyTorch就像一把瑞士军刀——你得先知道刀刃在哪，才能切菜。这一章，咱们就把PyTorch最核心的四个基本功打扎实。
2.1 张量操作：PyTorch的“数字积木”
张量是什么？说白了就是多维数组。标量是0维，向量是1维，矩阵是2维，再往上堆就是高维张量。我在项目中遇到过最典型的场景：图像数据进来是[batch, channel, height, width]的四维张量，你如果对维度操作不熟，后面reshape、permute的时候分分钟想砸键盘。
创建张量的几种姿势
import torch
# 从列表创建
a = torch.tensor([[1, 2], [3, 4]])
# 全零/全一
b = torch.zeros(3, 4)
c = torch.ones(2, 3)
# 随机初始化
d = torch.randn(2, 2) # 标准正态分布
# 指定设备
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
e = torch.randn(2, 2, device=device)
dtype。默认是float32，但如果你做量化或混合精度训练，float16能省一半显存。我曾经因为忘记设dtype，一个模型在V100上直接OOM，排查了半天才发现是默认float32在作祟。
张量运算与广播机制
加减乘除这些基础运算我就不啰嗦了。重点说说广播——这是PyTorch里最容易踩坑的地方之一。
# 广播示例
x = torch.randn(3, 1) # 形状 [3, 1]
y = torch.randn(1, 4) # 形状 [1, 4]
z = x + y # 结果形状 [3, 4]
为什么会这样？PyTorch会自动把维度为1的那一列“复制”成跟对方匹配的大小。但如果你不小心把形状搞成[3, 4]和[4, 3]，它不会报错，而是给你一个莫名其妙的结果。嗯，这里要注意：广播是从最后一个维度开始匹配的。
2.2 自动求导：反向传播的“自动驾驶”
我记得刚学深度学习那会儿，手动推导梯度公式写到崩溃。PyTorch的autograd模块就是来救命的。你只需要把张量的requires_grad设为True，它就会自动记录所有操作，然后调用.backward()就能拿到梯度。
x = torch.tensor([2.0], requires_grad=True)
y = x ** 2 + 3 * x + 1
y.backward()
print(x.grad) # 输出: tensor([7.]) 因为 dy/dx = 2x + 3 = 7
backward()之前，一定要调用optimizer.zero_grad()。这是新手最容易犯的错误，没有之一。
还有一个细节：with torch.no_grad()上下文管理器。在模型推理或计算验证集loss时，记得用这个包裹起来。它不会追踪梯度，能省显存还提速。你想想看，推理阶段你根本不需要梯度，何必让PyTorch白干活？
2.3 构建简单全连接网络
全连接网络，也叫多层感知机（MLP）。结构很简单：输入层 -> 隐藏层 -> 输出层，层与层之间全连接。咱们用nn.Module来定义，这是PyTorch里所有网络的基类。
import torch.nn as nn
class SimpleMLP(nn.Module):
def __init__(self, input_dim, hidden_dim, output_dim):
super().__init__()
self.fc1 = nn.Linear(input_dim, hidden_dim)
self.relu = nn.ReLU()
self.fc2 = nn.Linear(hidden_dim, output_dim)
def forward(self, x):
x = self.fc1(x)
x = self.relu(x)
x = self.fc2(x)
return x
# 实例化
model = SimpleMLP(784, 256, 10)
print(model)
我个人习惯把每一层都写在__init__里，forward只负责数据流。这样代码清晰，调试也方便。如果你用nn.Sequential，虽然写起来快，但想打印中间层的输出就麻烦了。
2.4 数据加载器与训练循环
数据加载器DataLoader是PyTorch里被低估的神器。它帮你做三件事：打乱数据、分批次、多进程加载。我建议你永远不要手动写循环来切batch，直接用DataLoader，省心又高效。
from torch.utils.data import DataLoader, TensorDataset
# 模拟数据
X = torch.randn(1000, 784)
y = torch.randint(0, 10, (1000,))
dataset = TensorDataset(X, y)
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)
# 训练循环
model = SimpleMLP(784, 256, 10)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
num_epochs = 10
for epoch in range(num_epochs):
running_loss = 0.0
for batch_x, batch_y in dataloader:
# 前向传播
outputs = model(batch_x)
loss = criterion(outputs, batch_y)
# 反向传播
optimizer.zero_grad()
loss.backward()
optimizer.step()
running_loss += loss.item()
print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {running_loss/len(dataloader):.4f}')
这个循环模板，你以后写任何PyTorch项目都会用到。我建议你把它背下来，或者存成代码片段。注意loss.item()——它把张量转成Python标量，不然你会得到一个张量列表，打印出来全是tensor(0.1234)，看着难受。
batch_x = batch_x.to(device)把数据搬过去。我见过有人忘了这一步，结果模型在GPU上算，数据在CPU上，PyTorch自动帮你搬，但每次迭代都搬一次，速度慢得离谱。
2.5 模型保存与加载
训练一个模型可能要几个小时甚至几天，你总不想每次重启都重头来吧？PyTorch提供了两种保存方式：保存整个模型，或者只保存状态字典。我强烈建议你只保存state_dict，原因有二：一是文件小，二是兼容性好——换了PyTorch版本也不容易报错。
# 保存
torch.save(model.state_dict(), 'model_weights.pth')
# 加载
model = SimpleMLP(784, 256, 10)
model.load_state_dict(torch.load('model_weights.pth'))
model.eval() # 切换到评估模式
model.eval()。如果你模型里有Dropout或BatchNorm，训练模式和评估模式的行为完全不同。不切换的话，推理结果会随机变化，排查起来特别痛苦。记住：训练用model.train()，推理用model.eval()。
如果你要保存checkpoint以便恢复训练，建议把epoch、optimizer状态、loss都存下来：
checkpoint = {
'epoch': epoch,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'loss': loss,
}
torch.save(checkpoint, 'checkpoint.pth')
这样即使训练到一半断电了，你也能从断点处继续，不用从头再来。
本章知识体系
下面这张图把本章的核心逻辑串起来了。你可以把它当作一张“地图”，学完后再回来看，应该会有更清晰的认识。
从张量操作开始，到自动求导帮你算梯度，再到用nn.Module搭网络，最后用DataLoader和训练循环把整个流程跑通——每一步都环环相扣。模型保存与加载则是你辛苦训练的“保险柜”。
这一章的内容，说白了就是PyTorch的“四件套”。你把这四个点吃透了，后面学CNN、RNN、Transformer都会轻松很多。别急，慢慢来，代码多敲几遍，手感自然就有了。
公众号:蓝海资料掘金营，微信deep3321