
PyTorch学习笔记小土堆这份笔记是我跟着小土堆的PyTorch系列视频一边看一边敲代码整理出来的。说是笔记其实更像是我从零开始跑通深度学习的一整套踩坑记录——从安装环境、张量操作到自定义数据集、搭网络、训练循环、模型保存每一个能出错的环节我都替你们错过了。PyTorch学习笔记这个关键词对应的内容网上很多但我尽量不写成文档翻译而是写成我当时是怎么想的、为什么这么写、换了你会不会踩坑。适合的读者很明确懂一点Python、见过神经网络名词但没真正跑过一个项目的同学。如果你已经会训练模型了这份笔记也可以当索引用里面有几处错误排查和多进程相关的经验属于官方教程里不太会写、但实际一定会碰到的东西。1.1 这份笔记的目标我要解决的问题只有一个让你用一口贴地的中文彻底搞清楚PyTorch从数据到模型的完整流程。不要纠结高阶API不要背复杂论文先把训练一个分类模型的最小闭环画出来。这个闭环包括四段数据怎么进来Dataset/DataLoader、网络怎么搭nn.Module、参数怎么更新loss.backward optimizer.step、模型怎么活下来save/load。笔记的主线就是这四件事其余的张量、transforms、设备切换都是为它们服务的。2. 环境搭建版本匹配和第一个张量很多人学PyTorch死在了第一步——不是代码不会写而是conda虚拟环境里装了个CPU版训练时发现慢到怀疑人生或者更惨装完直接报找不到torch。这里我把自己最终稳定用的方案写出来。2.1 先想清楚CPU还是GPU说实话如果只是学语法、跑小模型CPU版完全够用。但一旦你开始跑卷积网络、处理图片数据集CPU版的训练速度会让你失去耐心。所以我建议从一开始就装GPU版哪怕你现在没有NVIDIA显卡也按GPU版的命令装以后换机器不会懵。判断环境的核心命令就三行nvidia-smi # 看显卡驱动和最高支持的CUDA版本 python -c import torch; print(torch.__version__) python -c import torch; print(torch.cuda.is_available())第三行输出True说明PyTorch成功用上了CUDA。这里有个反直觉的地方CUDA不是你在系统里单独装的那个CUDA Toolkit决定的而是PyTorch包编译时自带的CUDA运行库决定的。所以经常出现系统CUDA是11.8PyTorch是12.1也能跑的情况。别花时间纠结系统CUDA版本只要nvidia-smi显示的驱动版本够新随便装PyTorch官方推荐版本基本都行。2.2 conda还是pip的一己之见我两种都用过。conda的好处是能顺便解决一些非Python依赖但PyTorch的包太大conda经常在解析依赖上卡很久。现在我更推荐pip因为PyTorch官方很早就提供了清晰的pip安装索引地址跟版本选择器是同步的。如果你用的是Anaconda先建一个独立环境别把PyTorch装进base环境否则后面装OpenCV、scikit-learn的时候容易互相污染。conda create -n torchlearn python3.11 -y conda activate torchlearn pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121torchvision建议顺手装上它不只是图像模型库很多经典数据集CIFAR-10、MNIST和常用transformsToTensor、Normalize都直接从它导入写学习笔记离不开它。2.3 第一个张量验证环境环境配好后的hello world不是print字符串而是确认张量能搬上GPUimport torch x torch.arange(12).reshape(3, 4) print(x) print(x.device) # cpu if torch.cuda.is_available(): gpu_x x.to(cuda) print(gpu_x.device) # cuda:0顺便说一句torch.arange(12).reshape(3, 4)比torch.arange(12).view(3, 4)更适合新手因为reshape对内存不连续的张量也能处理view遇到不连续的情况会直接报错。我在没有仔细想清楚之前被view坑过好几次。3. 张量操作把NumPy思维升级成Tensor思维张量是PyTorch的地基它的操作看起来像NumPy但多了两个关键东西设备属性和梯度属性。这一章我按创建、索引、梯度三条线展开。3.1 创建张量的几种方式以及tensor和Tensor的差别很多人一开始分不清torch.tensor和torch.Tensor记住这个就够了小写的torch.tensor是根据你传入的数据推断数据类型的构造函数大写的torch.Tensor是类等价于torch.FloatTensor。a torch.tensor([1, 2, 3]) # dtypetorch.int64因为传的是整数 b torch.Tensor([1, 2, 3]) # dtypetorch.float32直接定位FloatTensor c torch.randn(3, 4) # 标准正态分布最常用 d torch.zeros(2, 5, dtypetorch.float32) e torch.ones_like(c) # 和c形状一样值为1实际写代码时randn随机初始化权重、zeros初始化偏置、arange造等差序列做索引用得多。还有一个常用函数是torch.full((2, 3), 7)生成全填7的固定值张量看起来鸡肋但做注意力掩码、paint-padding的时候很顺手。3.2 索引、切片和广播张量的索引规则和NumPy几乎一模一样x[0]取第一行x[:, 1]取第二列x[1:, :2]是切片。真正需要动脑的是广播机制——形状不一致时低维度的张量会自动扩展。a torch.ones(3, 1) b torch.ones(1, 4) c a b # 结果是3x4a和b各自广播这个机制的坑在于它不会修正维度只是补维度。比如(3, 1)和(1, 4)能广播成(3, 4)但(3, 4)和(4,)能不能加能但结果是把(4,)当成行向量和每一行对齐相加。如果不理解这一点写损失函数、写注意力矩阵分分钟出现形状莫名其妙对上了但又不对的bug。建议遇到形状不对先打印两个张量的shape再想想广播规则别硬调。3.3 requires_grad张量最独特的脾气PyTorch张量和NumPy数组最本质的区别就是可以携带计算图并自动求梯度。打开这个功能的开关是requires_gradTrue但要注意只有浮点类型张量才能开这个开关整数张量直接报错。x torch.tensor([2.0], requires_gradTrue) y x ** 2 3 * x y.backward() print(x.grad) # tensor([7.]) 也就是2*x3在x2时的值几个经验性要点每次backward()默认是累加梯度的。想在循环里更新参数必须用optimizer.zero_grad()清零否则梯度会把之前的加在一起导致参数更新震荡。用x.detach()可以把张量从计算图中拿下来得到的张量不再需要梯度常用于记录指标、画图。在torch.no_grad()上下文里做验证和保存模型既省显存又不会污染计算图。至于自动求导的数学原理笔记阶段不用深究你只需要知道任何标量在对requires_grad的张量做了运算后调用backward()每个叶子节点的.grad就会被填上对应的偏导数。4. 数据管线Dataset、transforms、DataLoader三者配合跑深度学习时喂给模型的不是一张张图片而是一批批张量。PyTorch把这件事拆成了三层Dataset负责找到并读取一条样本transforms负责把样本处理成模型能吃的数据DataLoader负责按batch打包并乱序。这个设计初看多此一举等你自己写数据预处理时就会感谢它。4.1 自定义Dataset的万能写法官方教程里的Dataset示例非常简单但现实中的数据集动不动就是文件夹分类、CSV带label、路径列表。我建议直接背住这个模板from torch.utils.data import Dataset from PIL import Image import os class MyImageDataset(Dataset): def __init__(self, image_dir, transformNone): self.image_paths [] self.labels [] for label, folder in enumerate(os.listdir(image_dir)): folder_path os.path.join(image_dir, folder) for fname in os.listdir(folder_path): self.image_paths.append(os.path.join(folder_path, fname)) self.labels.append(label) self.transform transform def __len__(self): return len(self.image_paths) def __getitem__(self, idx): img Image.open(self.image_paths[idx]).convert(RGB) label self.labels[idx] if self.transform: img self.transform(img) return img, label__len__返回样本总数__getitem__根据索引返回一条数据。DataLoader内部就是反复调用__getitem__来凑batch的。如果你用open打开图片时遇到奇怪的颜色通道错误convert(RGB)能统一通道数这是最常见的一个坑。4.2 transforms到底在做什么从文件读进来的是PIL图片是HWC、数值范围0-255、数据类型是uint8。神经网络要的是CHW、范围0-1或者标准化之后、float32。transforms就是干这件事的。from torchvision import transforms transform transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ])关键理解Compose里的每个变换都是按顺序执行的前一个的输出是后一个的输入。Resize必须在ToTensor之前因为ToTensor之后数据类型变了Normalize必须在ToTensor之后因为它按通道做标准化需要float类型。如果你把顺序搞反报错信息会很长但本质就俩字类型。4.3 DataLoader参数的含义和性能关系from torch.utils.data import DataLoader loader DataLoader( dataset, batch_size32, shuffleTrue, num_workers4, drop_lastFalse, pin_memoryTrue )shuffleTrue只对训练集开验证集必须False否则指标失去意义。drop_lastTrue会丢掉最后一批不足batch_size的数据。训练一般开True因为最后一批太小会让BN和梯度估计抖动验证一般False。num_workers是加载数据的子进程数。在GPU训练时数据加载往往成为瓶颈调大它能明显提速。但如果你在Windows环境跑把它设成非0前必须把主代码包在if __name__ __main__:里否则会递归启动子进程导致程序崩溃。这个错误我在初学阶段反复踩后来养成了固定写这个判断语句的习惯。pin_memoryTrue配合GPU训练能省一步数据拷贝建议常开。5. 网络搭建nn.Module其实是个容器很多新手把nn.Module当成必须继承的类然后就开始背模板。我更愿意把它理解成一个容器你在__init__里声明层在forward里定义数据的流向整个容器自动获得参数管理、train/eval切换、设备迁移的能力。5.1 nn.Module的基本写法和执行顺序import torch.nn as nn import torch.nn.functional as F class SimpleNet(nn.Module): def __init__(self): super().__init__() self.fc1 nn.Linear(784, 128) self.fc2 nn.Linear(128, 10) def forward(self, x): x x.view(x.size(0), -1) # 展平 x F.relu(self.fc1(x)) x self.fc2(x) return x这里的顺序有讲究先展平再全连接再接激活函数。F.relu是函数式接口用完就释放nn.ReLU是模块接口会有自己的存储。在小网络里用哪个都行但记住一个原则__init__里只管有参数的层无参数的激活函数、dropout、flatten尽量写在forward里用函数式接口这样模型文件更轻、结构更清楚。5.2 常用层的参数细节nn.Conv2d(in_channels, out_channels, kernel_size, stride1, padding0)。padding不是必须填0填1就是周围加一圈0。output的尺寸公式是(H - kernel 2*padding) / stride 1。我建议你至少自己推导一次7x7输入、3x3核、padding1后的尺寸以后搭网络不用每次都算。nn.MaxPool2d(kernel_size, stride)默认stride等于kernel_size很多人的参数写法就省略了这不一定会错但建议显式写出来。nn.BatchNorm2d(num_features)传的是通道数不是别的。它在batch上做归一化训练时用本批次统计量eval时用滑动平均统计量。这就引出一个经典坑验证时忘记model.eval()结果指标一直在变。nn.Dropout(p)p是丢弃比例p0.5表示一半神经元失活。它只在训练时生效eval时自动失效。5.3 搭建小型卷积网络的完整示范以CIFAR-10的三通道32x32图片为例这是我笔记里用的第一个稍微像样的网络class TinyCNN(nn.Module): def __init__(self): super().__init__() self.conv1 nn.Conv2d(3, 16, 3, padding1) # 32x32 - 32x32 self.conv2 nn.Conv2d(16, 32, 3, padding1) # 32x32 - 32x32 self.pool nn.MaxPool2d(2, 2) # 32x32 - 16x16 self.fc1 nn.Linear(32 * 16 * 16, 128) self.fc2 nn.Linear(128, 10) def forward(self, x): x F.relu(self.conv1(x)) x self.pool(F.relu(self.conv2(x))) x x.view(x.size(0), -1) x F.relu(self.fc1(x)) x self.fc2(x) return x这里唯一的计算难点是fc1的输入维度。如果你记不住直接打印一下x.shape就行但我强烈建议你手动算一次输入32x32第一个卷积padding1核3x3尺寸不变池化核2x2尺寸减半到16x16通道数是32。所以展平后是32*16*16。学会这种推尺寸的方法之后任何网络结构你都能看懂参数从哪来。6. 训练闭环损失、优化器、反向传播的执行顺序我见过不少同学能把网络搭起来、数据跑起来但一到训练循环就乱了zero_grad放哪、backward放哪、step放哪顺序错了还没报错但loss就是降不下去。这一章把最完整的循环写出来然后逐个解释为什么是这个顺序。6.1 一个完整的训练循环import torch import torch.nn as nn import torch.optim as optim model TinyCNN() criterion nn.CrossEntropyLoss() optimizer optim.Adam(model.parameters(), lr1e-3) num_epochs 10 for epoch in range(num_epochs): model.train() total_loss 0.0 for images, labels in train_loader: # 1. 梯度清零这是最容易漏的一步 optimizer.zero_grad() # 2. 前向传播 outputs model(images) # 3. 计算损失 loss criterion(outputs, labels) # 4. 反向传播 loss.backward() # 5. 更新参数 optimizer.step() total_loss loss.item() avg_loss total_loss / len(train_loader) print(fEpoch [{epoch1}/{num_epochs}], Loss: {avg_loss:.4f})model.train()这行看起来很废话但它决定Dropout和BatchNorm的行为。很多人在训练时忘了写在eval时忘了切回train导致同一个模型两次结果不一样。我后来养成的习惯是训练循环开头写model.train()验证循环开头写model.eval()雷打不动。6.2 为什么顺序是zero_grad、backward、step分开来解释zero_grad()是把上一次迭代留下的梯度清空。PyTorch的梯度是累加的不清空的话第二次的梯度会叠加到第一次上这时候的梯度是错的参数更新也会错。我当初天真地以为backward会自动覆盖梯度结果loss曲线上下乱跳查了半天才发现是这一个零函数的问题。backward()做的是反向传播它的效果是从loss出发把梯度按计算图反推填到每个requires_gradTrue的叶子节点也就是model.parameters()里的那些参数对应的.grad上。这步执行完之后计算图默认会被释放所以如果你的程序里想对同一批数据做两次backward会被提示中间变量被释放了。step()才是真正的参数更新optimizer根据.grad去修改param.data。所以三者顺序必须是zero_grad - backward - step。如果你把zero_grad放在step之后等于上一轮的梯度参与了这一轮的参数更新如果你忘了zero_grad等于每一轮都拿着历史梯度的大杂烩。6.3 训练中常见的三个问题丢失不降先看学习率如果loss从头到尾几乎不变大概率lr太小如果loss直接变NaN大概率lr太大。然后看数据预处理标签是否从0开始连续编号、特征是否归一化。我的检查顺序是先跑一批数据过一遍forward确认输出维度对再单独跑一次backward确认没有维度错最后才检查学习率。梯度爆炸或NaN除了学习率过大还有可能是输入包含NaN值。检查方法是在loss计算前打印torch.isnan(images).sum()。对于文本、图像类数据这个概率低但自定义数据集里很容易出现脏数据。训练慢数据加载是主要瓶颈。调大num_workers、给数据加pin_memoryTrue多半能解决。还有一个不起眼的点GPU训练时把input和label也搬到设备上时尽量一次images, label images.to(device), label.to(device)完成别反复跨越设备。7. 模型持久化保存、加载、切换设备的现场经验模型训练完了要保存推理时再加载。这个环节看起来简单但藏着不少细节直接torch.save(model)虽然能存但换环境就崩只存参数又不建网络一样加载不了。这一章我把经过验证的方案写全。7.1 state_dict和整体保存的区别PyTorch官方推荐保存的是model.state_dict()它是一个Python字典把每个有参数层的名字映射到对应的张量上。torch.save(model.state_dict(), tiny_cnn.pth) # 加载时先实例化一个结构完全相同的模型再load loaded_model TinyCNN() loaded_model.load_state_dict(torch.load(tiny_cnn.pth))torch.save(model, xxx.pth)这种整体保存看着方便但它会把网络结构定义也序列化进去。问题出在如果你把脚本复制到另一台机器或者类的定义路径变了就可能反序列化失败报无法找到类之类的错。所以在正式场合我只推荐存state_dict。还有一个实用的扩展把训练时的重要超参数和数据预处理配置也放进同一个字典存成一个.pth文件避免以后看模型文件时还要翻代码回忆当时的输入尺寸和mean/std。7.2 加载后必须做的两件事loaded_model.load_state_dict(torch.load(tiny_cnn.pth)) loaded_model.eval()第一件事load_state_dict之后模型默认处于train模式直接拿去推理Dropout会随机丢弃神经元BatchNorm也会用训练时统计量之外的东西输出结果是错乱的。所以必须显式调用eval()。第二件事如果加载的模型原来在GPU上训练保存的参数也在GPU上而你目前跑在CPU机器torch.load会报设备不匹配的错误。解法是加载时指定map_locationstate torch.load(tiny_cnn.pth, map_locationcpu) loaded_model.load_state_dict(state)7.3 CPU/GPU推理切换的优雅写法我在笔记最后写了一个比较通用的推理片段无论跑在哪都行device cuda if torch.cuda.is_available() else cpu model.to(device) def predict(image_tensor): model.eval() with torch.no_grad(): image image_tensor.unsqueeze(0).to(device) # 加batch维 logits model(image) pred logits.argmax(dim1).item() return predunsqueeze(0)是常被漏掉的一步模型接受的是带batch维的输入单张图片预测时必须先增加一个维度。在推理代码里包上torch.no_grad()能省显存、加快速度还不会误触发计算图的构造。这两点我在早期都吃过亏一个让模型直接报维度错误一个让推理慢了一倍还额外耗显存。8. 整理这份笔记时反复踩过的坑最后这部分是散装经验不是系统教程但每一条都是我亲手踩出来的拿出来给大家避雷。第一个坑是把.item()挂在所有张量后面。loss.item()很好用它是从只含一个元素的张量里取Python标量。但如果你对logits这种多维张量也调.item()会直接报错。正确做法是loss.item()、pred.argmax(dim1).item()这类明确得到标量的场景才用。想同时看多个元素的数值用.tolist()更稳。第二个坑是设备不一致引发的hidden device报错。参数在cuda、输入张量在cpu一forward就报错。排查方法很简单在forward第一行打印x.device和self.fc1.weight.device一眼就能看出谁没搬。这个报错的英文提示很模糊新手最容易在这卡半小时。第三个坑是克隆与原地修改。如果你写了类似params model.fc1.weight然后对这个引用做原地操作比如切片赋值或fill_会连带改变原本的参数。需要用weight.clone()做一份真正独立的副本。类似地对输入张量做img - mean这类原地操作时要确认没有和原图引用纠缠否则数据会被悄悄改掉。第四个坑是DataLoader的shuffle顺序。验证集千万不能shuffle否则每个epoch的样本顺序都变测试指标即使模型没变也会轻微抖动排查起来非常迷惑。我一度怀疑模型没收敛后来发现只是验证集的shuffle没关。最后一个是一些比较隐蔽的性能问题如果模型里用了多种dtype混着算比如把float32参数和float64输入相加PyTorch会先统一类型白白多花一次转换时间。我自己写代码时经常在Dataset的__getitem__返回前统一做images images.float()保证进模型的张量类型一致。这份笔记写到这里其实已经把PyTorch从数据到模型的最小闭环完整走了一遍。学完这些内容如果你再去读官方文档的60分钟入门教程会发现自己已经能直接跳过一大半内容了。至于接下来是去啃nn.Transformer还是自己实现一个注意力机制那就是以后的事了——先把这份笔记里的代码亲手敲一遍比什么都强。