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

📅 2026/7/5 6:39:40
从零手搓大模型前置知识(附录一)PyTorch 基础,从张量到训练循环
从零手搓大模型前置知识附录一PyTorch 基础从张量到训练循环这一部分是本系列的 PyTorch 入门补充。它不是直接讲大模型而是帮你补齐后面手搓 LLM 必须用到的 PyTorch 基础tensor - 自动求导 - 神经网络模块 - Dataset/DataLoader - 训练循环 - 保存和加载模型如果你看从零手搓大模型系列正文时对torch.tensor、nn.Module、DataLoader、loss.backward()感到陌生建议先把这篇学完。1. PyTorch 是什么PyTorch 是一个深度学习框架。对我们来说它主要提供三类能力用 tensor 表示数据和模型参数。自动计算梯度。用 GPU 加速训练。导入 PyTorchimporttorch检查 GPU 是否可用print(torch.cuda.is_available())如果输出True说明当前环境可以使用 NVIDIA GPU。输出False也没关系附录一和前几章代码大多可以在 CPU 上跑。2. TensorPyTorch 的基本数据结构tensor 可以理解成“支持自动求导和 GPU 加速的多维数组”。常见形式0D tensor: scalar标量 1D tensor: vector向量 2D tensor: matrix矩阵 3D tensor: 更高维张量示例importtorchimportnumpyasnp tensor0dtorch.tensor(1)tensor1dtorch.tensor([1,2,3])tensor2dtorch.tensor([[1,2],[3,4]])tensor3d_1torch.tensor([[[1,2],[3,4]],[[5,6],[7,8]]])这些就是后面 LLM 里所有数据的基础。比如第1章里token IDs: 2D tensor token embeddings: 3D tensor第2章里attention scores: 2D 或 4D tensor3. 从 NumPy 转成 tensorNumPy 和 PyTorch 的关系ary3dnp.array([[[1,2],[3,4]],[[5,6],[7,8]]])tensor3d_2torch.tensor(ary3d)tensor3d_3torch.from_numpy(ary3d)区别很重要torch.tensor(ary): 会复制一份数据 torch.from_numpy(ary): 和 NumPy 数组共享内存共享内存意味着如果你改了原来的 NumPy 数组对应的 tensor 也可能跟着变。初学时记住一个简单建议想要安全复制用 torch.tensor(...) 想要省内存共享用 torch.from_numpy(...)4. Tensor 的数据类型 dtype示例tensor1dtorch.tensor([1,2,3])print(tensor1d.dtype)整数列表默认通常会得到torch.int64浮点列表floatvectorch.tensor([1.0,2.0,3.0])通常会得到torch.float32也可以手动转换floatvectensor1d.to(torch.float32)为什么 dtype 重要因为神经网络的权重和输入通常是浮点数例如float32、float16、bfloat16。而分类标签常常是整数例如int64。5. 常见 tensor 操作创建一个二维 tensortensor2dtorch.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_scoresqueries keys.T所以一定要熟悉它就是矩阵乘法。6. 把模型看成计算图用一个最小神经元示例importtorch.nn.functionalasF ytorch.tensor([1.0])# true labelx1torch.tensor([1.1])# input featurew1torch.tensor([2.2])# weight parameterbtorch.tensor([0.0])# bias unitzx1*w1b atorch.sigmoid(z)lossF.binary_cross_entropy(a,y)print(loss)流程是输入 x1 - 乘权重 w1加偏置 b - 得到 z - sigmoid 激活得到 a - 和真实标签 y 计算 loss这就是一个计算图。神经网络再复杂本质也是很多这样的计算节点连起来。7. 自动求导 autograd训练神经网络时我们要知道loss 对每个参数的梯度是多少PyTorch 可以自动计算。示例importtorch.nn.functionalasFfromtorch.autogradimportgrad ytorch.tensor([1.0])x1torch.tensor([1.1])w1torch.tensor([2.2],requires_gradTrue)btorch.tensor([0.0],requires_gradTrue)zx1*w1b atorch.sigmoid(z)lossF.binary_cross_entropy(a,y)grad_L_w1grad(loss,w1,retain_graphTrue)grad_L_bgrad(loss,b,retain_graphTrue)print(grad_L_w1)print(grad_L_b)关键是requires_gradTrue它告诉 PyTorch这个 tensor 是需要训练的参数请记录它参与的计算并能对它求梯度。后面训练循环里更常见的是loss.backward()它会自动沿着计算图反向传播把所有可训练参数的梯度算出来。8. 用 nn.Module 定义神经网络定义一个小型多层神经网络classNeuralNetwork(torch.nn.Module):def__init__(self,num_inputs,num_outputs):super().__init__()self.layerstorch.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):logitsself.layers(x)returnlogits几个重点torch.nn.Module是所有 PyTorch 模型的基类。__init__里定义层torch.nn.Linear(...)torch.nn.ReLU()forward里定义数据怎么流过模型logitsself.layers(x)这和后面 GPT 模型是同一套路定义模块 - 写 forward - 输入 tensor - 输出 logits9. logits 是什么模型最后输出logitsmodel(x)logits 是还没有经过 softmax 的原始分数。比如二分类时可能输出[2.1, -0.8]第一个类别分数更高就预测类别 0。训练时F.cross_entropy可以直接吃 logits不需要你先手动 softmax。10. Dataset定义数据集notebook 先准备玩具训练数据X_traintorch.tensor([[-1.2,3.1],[-0.9,2.9],[-0.5,2.6],[2.3,-1.1],[2.7,-1.5]])y_traintorch.tensor([0,0,0,1,1])然后定义 Datasetfromtorch.utils.dataimportDatasetclassToyDataset(Dataset):def__init__(self,X,y):self.featuresX self.labelsydef__getitem__(self,index):one_xself.features[index]one_yself.labels[index]returnone_x,one_ydef__len__(self):returnself.labels.shape[0]Dataset 必须实现两个方法__len__: 返回数据集大小 __getitem__: 根据 index 返回一条样本第1章里的GPTDatasetV1也是同样思想只不过返回的是input_ids, target_ids11. DataLoader批量取数据定义 DataLoaderfromtorch.utils.dataimportDataLoader torch.manual_seed(123)train_loaderDataLoader(datasettrain_ds,batch_size2,shuffleTrue,num_workers0)参数解释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)modelNeuralNetwork(num_inputs2,num_outputs2)optimizertorch.optim.SGD(model.parameters(),lr0.5)num_epochs3forepochinrange(num_epochs):model.train()forbatch_idx,(features,labels)inenumerate(train_loader):logitsmodel(features)lossF.cross_entropy(logits,labels)optimizer.zero_grad()loss.backward()optimizer.step()print(fEpoch:{epoch1:03d}/{num_epochs:03d}f | Batch{batch_idx1: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后要stepbackward只负责算梯度step才真正更新参数。zero_grad()只负责清理历史垃圾梯度不产生新梯度loss.backward()才是生成当前批次梯度的唯一步骤13. model.train() 和 model.eval()训练时model.train()评估时model.eval()它们会影响 dropout、batch norm 等层的行为。虽然这个小模型里影响不明显但后面大模型里一定要养成习惯。14. 计算准确率定义defcompute_accuracy(model,dataloader):model.eval()correct0.0total_examples0foridx,(features,labels)inenumerate(dataloader):withtorch.no_grad():logitsmodel(features)predictionstorch.argmax(logits,dim1)comparelabelspredictions correcttorch.sum(compare)total_exampleslen(compare)return(correct/total_examples).item()关键点withtorch.no_grad():评估时不需要求梯度可以省内存、省计算。预测类别predictionstorch.argmax(logits,dim1)意思是对每个样本取 logits 最大的类别。15. 保存和加载模型保存torch.save(model.state_dict(),model.pth)加载modelNeuralNetwork(2,2)model.load_state_dict(torch.load(model.pth,weights_onlyTrue))注意加载时模型结构必须和保存时一样。state_dict保存的是参数不是完整 Python 类定义。所以你要先创建同样结构的模型modelNeuralNetwork(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. 建议按这个顺序学先跑 tensor 部分熟悉 shape、dtype、reshape、transpose、矩阵乘法。再跑自动求导理解requires_gradTrue和loss.backward()。然后看NeuralNetwork理解nn.Module和forward。再看 Dataset/DataLoader理解数据怎么按 batch 送进模型。最后认真看训练循环把那 6 步背熟。如果只记一个训练模板记这个model.train()forfeatures,labelsintrain_loader:logitsmodel(features)lossF.cross_entropy(logits,labels)optimizer.zero_grad()loss.backward()optimizer.step()后面训练 GPT本质上也是这个套路只是模型更大、数据更复杂。