从零手搓大模型前置知识（附录一）PyTorch 基础，从张量到训练循环

从零手搓大模型前置知识（附录一）PyTorch 基础，从张量到训练循环

这一部分是本系列的 PyTorch 入门补充。它不是直接讲大模型，而是帮你补齐后面手搓 LLM 必须用到的 PyTorch 基础：

tensor -> 自动求导 -> 神经网络模块 -> Dataset/DataLoader -> 训练循环 -> 保存和加载模型

如果你看从零手搓大模型系列正文时对torch.tensor、nn.Module、DataLoader、loss.backward()感到陌生，建议先把这篇学完。

1. PyTorch 是什么

PyTorch 是一个深度学习框架。对我们来说，它主要提供三类能力：

1. 用 tensor 表示数据和模型参数。
2. 自动计算梯度。
3. 用 GPU 加速训练。

导入 PyTorch：

importtorch

检查 GPU 是否可用：

print(torch.cuda.is_available())

如果输出True，说明当前环境可以使用 NVIDIA GPU。输出False也没关系，附录一和前几章代码大多可以在 CPU 上跑。

2. Tensor：PyTorch 的基本数据结构

tensor 可以理解成“支持自动求导和 GPU 加速的多维数组”。

常见形式：

0D tensor: scalar，标量 1D tensor: vector，向量 2D tensor: matrix，矩阵 3D+ tensor: 更高维张量

示例：

importtorchimportnumpyasnp tensor0d=torch.tensor(1)tensor1d=torch.tensor([1,2,3])tensor2d=torch.tensor([[1,2],[3,4]])tensor3d_1=torch.tensor([[[1,2],[3,4]],[[5,6],[7,8]]])

这些就是后面 LLM 里所有数据的基础。

比如第1章里：

token IDs: 2D tensor token embeddings: 3D tensor

第2章里：

attention scores: 2D 或 4D tensor

3. 从 NumPy 转成 tensor

NumPy 和 PyTorch 的关系：

ary3d=np.array([[[1,2],[3,4]],[[5,6],[7,8]]])tensor3d_2=torch.tensor(ary3d)tensor3d_3=torch.from_numpy(ary3d)

区别很重要：

torch.tensor(ary): 会复制一份数据 torch.from_numpy(ary): 和 NumPy 数组共享内存

共享内存意味着，如果你改了原来的 NumPy 数组，对应的 tensor 也可能跟着变。

初学时记住一个简单建议：

想要安全复制，用 torch.tensor(...) 想要省内存共享，用 torch.from_numpy(...)

4. Tensor 的数据类型 dtype

示例：

tensor1d=torch.tensor([1,2,3])print(tensor1d.dtype)

整数列表默认通常会得到：

torch.int64

浮点列表：

floatvec=torch.tensor([1.0,2.0,3.0])

通常会得到：

torch.float32

也可以手动转换：

floatvec=tensor1d.to(torch.float32)

为什么 dtype 重要？

因为神经网络的权重和输入通常是浮点数，例如float32、float16、bfloat16。而分类标签常常是整数，例如int64。

5. 常见 tensor 操作

创建一个二维 tensor：

tensor2d=torch.tensor([[1,2,3],[4,5,6]])

查看形状：

tensor2d.shape

这里形状是：

(2, 3)

表示 2 行 3 列。

改变形状：

tensor2d.reshape(3,2)

也可以用：

tensor2d.view(3,2)

初学时可以先把reshape和view都理解为“改形状”。更细的区别以后再看。

转置：

tensor2d.T

矩阵乘法：

tensor2d.matmul(tensor2d.T)

更常见写法：

tensor2d @ tensor2d.T

这在第2章 attention 里会大量出现：

attn_scores=queries @ keys.T

所以@一定要熟悉，它就是矩阵乘法。

6. 把模型看成计算图

用一个最小神经元示例：

importtorch.nn.functionalasF y=torch.tensor([1.0])# true labelx1=torch.tensor([1.1])# input featurew1=torch.tensor([2.2])# weight parameterb=torch.tensor([0.0])# bias unitz=x1*w1+b a=torch.sigmoid(z)loss=F.binary_cross_entropy(a,y)print(loss)

流程是：

输入 x1 -> 乘权重 w1，加偏置 b -> 得到 z -> sigmoid 激活得到 a -> 和真实标签 y 计算 loss

这就是一个计算图。

神经网络再复杂，本质也是很多这样的计算节点连起来。

7. 自动求导 autograd

训练神经网络时，我们要知道：

loss 对每个参数的梯度是多少

PyTorch 可以自动计算。

示例：

importtorch.nn.functionalasFfromtorch.autogradimportgrad y=torch.tensor([1.0])x1=torch.tensor([1.1])w1=torch.tensor([2.2],requires_grad=True)b=torch.tensor([0.0],requires_grad=True)z=x1*w1+b a=torch.sigmoid(z)loss=F.binary_cross_entropy(a,y)grad_L_w1=grad(loss,w1,retain_graph=True)grad_L_b=grad(loss,b,retain_graph=True)print(grad_L_w1)print(grad_L_b)

关键是：

requires_grad=True

它告诉 PyTorch：

这个 tensor 是需要训练的参数，请记录它参与的计算，并能对它求梯度。

后面训练循环里更常见的是：

loss.backward()

它会自动沿着计算图反向传播，把所有可训练参数的梯度算出来。

8. 用 nn.Module 定义神经网络

定义一个小型多层神经网络：

classNeuralNetwork(torch.nn.Module):def__init__(self,num_inputs,num_outputs):super().__init__()self.layers=torch.nn.Sequential(torch.nn.Linear(num_inputs,30),torch.nn.ReLU(),torch.nn.Linear(30,20),torch.nn.ReLU(),torch.nn.Linear(20,num_outputs),)defforward(self,x):logits=self.layers(x)returnlogits

几个重点：

torch.nn.Module是所有 PyTorch 模型的基类。

__init__里定义层：

torch.nn.Linear(...)torch.nn.ReLU()

forward里定义数据怎么流过模型：

logits=self.layers(x)

这和后面 GPT 模型是同一套路：

定义模块 -> 写 forward -> 输入 tensor -> 输出 logits

9. logits 是什么

模型最后输出：

logits=model(x)

logits 是还没有经过 softmax 的原始分数。

比如二分类时可能输出：

[2.1, -0.8]

第一个类别分数更高，就预测类别 0。

训练时F.cross_entropy可以直接吃 logits，不需要你先手动 softmax。

10. Dataset：定义数据集

notebook 先准备玩具训练数据：

X_train=torch.tensor([[-1.2,3.1],[-0.9,2.9],[-0.5,2.6],[2.3,-1.1],[2.7,-1.5]])y_train=torch.tensor([0,0,0,1,1])

然后定义 Dataset：

fromtorch.utils.dataimportDatasetclassToyDataset(Dataset):def__init__(self,X,y):self.features=X self.labels=ydef__getitem__(self,index):one_x=self.features[index]one_y=self.labels[index]returnone_x,one_ydef__len__(self):returnself.labels.shape[0]

Dataset 必须实现两个方法：

__len__: 返回数据集大小 __getitem__: 根据 index 返回一条样本

第1章里的GPTDatasetV1也是同样思想，只不过返回的是：

input_ids, target_ids

11. DataLoader：批量取数据

定义 DataLoader：

fromtorch.utils.dataimportDataLoader torch.manual_seed(123)train_loader=DataLoader(dataset=train_ds,batch_size=2,shuffle=True,num_workers=0)

参数解释：

• dataset：数据集对象。
• batch_size：每次取几条样本。
• shuffle：每轮训练是否打乱。
• num_workers：加载数据的子进程数，Windows 初学阶段用 0 更稳。

遍历：

foridx,(x,y)inenumerate(train_loader):print(idx,x,y)

训练模型时，我们不是一次只喂一条样本，而是一批一批喂。

这就是 batch training。

12. 典型训练循环

这是这部分最重要的代码：

importtorch.nn.functionalasF torch.manual_seed(123)model=NeuralNetwork(num_inputs=2,num_outputs=2)optimizer=torch.optim.SGD(model.parameters(),lr=0.5)num_epochs=3forepochinrange(num_epochs):model.train()forbatch_idx,(features,labels)inenumerate(train_loader):logits=model(features)loss=F.cross_entropy(logits,labels)optimizer.zero_grad()loss.backward()optimizer.step()print(f"Epoch:{epoch+1:03d}/{num_epochs:03d}"f" | Batch{batch_idx+1:03d}/{len(train_loader):03d}"f" | Train/Val Loss:{loss:.2f}")model.eval()

训练循环可以背成固定模板：

1. model.train() 2. 前向传播：logits = model(features) 3. 计算损失：loss = loss_fn(logits, labels) 4. 清空旧梯度：optimizer.zero_grad() 5. 反向传播：loss.backward() 6. 更新参数：optimizer.step() 7. model.eval()

为什么要zero_grad？

因为 PyTorch 默认会累积梯度。如果不清空，上一个 batch 的梯度会混进来。

为什么backward后要step？

backward只负责算梯度，step才真正更新参数。

zero_grad()只负责清理历史垃圾梯度，不产生新梯度；

loss.backward()才是生成当前批次梯度的唯一步骤；

13. model.train() 和 model.eval()

训练时：

model.train()

评估时：

model.eval()

它们会影响 dropout、batch norm 等层的行为。

虽然这个小模型里影响不明显，但后面大模型里一定要养成习惯。

14. 计算准确率

定义：

defcompute_accuracy(model,dataloader):model.eval()correct=0.0total_examples=0foridx,(features,labels)inenumerate(dataloader):withtorch.no_grad():logits=model(features)predictions=torch.argmax(logits,dim=1)compare=labels==predictions correct+=torch.sum(compare)total_examples+=len(compare)return(correct/total_examples).item()

关键点：

withtorch.no_grad():

评估时不需要求梯度，可以省内存、省计算。

预测类别：

predictions=torch.argmax(logits,dim=1)

意思是对每个样本，取 logits 最大的类别。

15. 保存和加载模型

保存：

torch.save(model.state_dict(),"model.pth")

加载：

model=NeuralNetwork(2,2)model.load_state_dict(torch.load("model.pth",weights_only=True))

注意：加载时模型结构必须和保存时一样。

state_dict保存的是参数，不是完整 Python 类定义。

所以你要先创建同样结构的模型：

model=NeuralNetwork(2,2)

再把权重加载进去。

16.本章 和 LLM 的关系

本章对后面章节的帮助非常直接：

第 1 章会用：

Dataset DataLoader tensor shape

第 2 章会用：

矩阵乘法 transpose softmax nn.Module

第 3 章会用：

nn.Module nn.Linear forward Sequential 思维

第 4 章训练 GPT 会完整用到：

loss.backward() optimizer.step() model.train() model.eval() torch.no_grad() state_dict

所以这部分不是“附录可看可不看”，而是后面手搓大模型的 PyTorch 地基。

17. 建议

按这个顺序学：

1. 先跑 tensor 部分，熟悉 shape、dtype、reshape、transpose、矩阵乘法。
2. 再跑自动求导，理解requires_grad=True和loss.backward()。
3. 然后看NeuralNetwork，理解nn.Module和forward。
4. 再看 Dataset/DataLoader，理解数据怎么按 batch 送进模型。
5. 最后认真看训练循环，把那 6 步背熟。

如果只记一个训练模板，记这个：

model.train()forfeatures,labelsintrain_loader:logits=model(features)loss=F.cross_entropy(logits,labels)optimizer.zero_grad()loss.backward()optimizer.step()

后面训练 GPT，本质上也是这个套路，只是模型更大、数据更复杂。