1. 从“Hello, Tensor”到第一个神经网络为什么PyTorch是2024年的首选如果你最近在任何一个技术社区或者招聘网站上扫一眼会发现一个高频出现的词PyTorch。无论是研究前沿大模型的实验室还是部署工业级视觉检测的生产线PyTorch几乎成了深度学习框架的代名词。这背后有一个非常直观的原因它让研究者、工程师甚至学生都能以一种更符合人类直觉的“Pythonic”方式去思考和构建模型。想象一下你写一个Python列表推导式或者操作一个NumPy数组那种流畅感PyTorch在构建和调试神经网络时提供了几乎一模一样的体验。我刚开始接触深度学习时也用过其他框架。那时候最痛苦的不是数学原理而是框架本身带来的心智负担——你需要先定义一个静态的“计算图”然后往里面“喂”数据调试起来像是在隔着一层毛玻璃看问题。PyTorch的出现彻底改变了这个局面。它的核心设计哲学是“动态计算图”和“命令式编程”这意味着你的代码执行顺序就是你的思维顺序。你可以像写普通Python程序一样在任意地方设置断点打印张量的值实时看到每一层输出的变化。这种“所见即所得”的编程模式极大地降低了入门和调试的门槛这也是为什么它能在学术界迅速流行并最终在工业界也占据主导地位。所以这篇教程的目标不是让你死记硬背API而是带你理解PyTorch之所以好用的“核心概念”。我们会从最基础的张量操作开始一步步搭建起对自动微分、计算图、优化器的认知最终亲手构建并训练一个能识别手写数字的神经网络。无论你是刚学完Python语法想进入AI领域的学生还是有一定其他框架经验想转过来的开发者这篇“快速入门”都会帮你建立一个坚实且正确的起点。记住我们的目标是理解原理掌握工具然后去创造。2. 环境搭建避开版本地狱一步到位配置GPU环境在开始写第一行代码之前一个正确且高效的环境是成功的一半。网络上大量的“PyTorch安装失败”求助帖根源大多在于版本不匹配——Python版本、PyTorch版本、CUDA版本、显卡驱动版本这四者必须形成一个兼容链。我们今天就彻底理清这个关系并给出一个2024年通用的、稳妥的安装方案。2.1 核心依赖关系解读CUDA、驱动与PyTorch的“三角恋”首先你必须理解这个依赖链条你的深度学习代码 - PyTorch库 - CUDA工具包 - NVIDIA显卡驱动 - 物理GPU硬件。物理GPU与驱动这是最底层。你需要一块NVIDIA显卡如RTX 3060, 4090等并安装最新的官方驱动。驱动版本决定了你的GPU能支持的最高CUDA版本。CUDA工具包这是NVIDIA推出的并行计算平台。PyTorch需要调用CUDA的库函数来在GPU上执行计算。你不需要单独完整安装CUDAPyTorch的预编译包已经自带了对应版本的CUDA运行时库。你只需要确保系统驱动支持该CUDA版本即可。PyTorch版本这是我们需要安装的库。在安装时你需要选择与你的目标CUDA版本匹配的PyTorch安装命令。如何查询自己的驱动支持的最高CUDA版本打开命令行Windows的CMD或PowerShellLinux/Mac的终端输入nvidia-smi命令。在输出结果的右上角你会看到类似CUDA Version: 12.4的字样。这个不是你安装的CUDA版本而是你的显卡驱动所能支持的最高CUDA版本。例如显示12.4意味着你可以安装需要CUDA 12.1, 11.8等只要≤12.4的PyTorch版本。2.2 实战安装使用Conda打造独立环境强烈建议使用Anaconda或Miniconda来管理环境。它可以为每个项目创建独立的Python环境避免包冲突。假设你已经安装了Conda。步骤一创建并激活新环境我们创建一个名为pytorch_tutorial的Python 3.9环境3.9是一个兼容性极好的版本。conda create -n pytorch_tutorial python3.9 conda activate pytorch_tutorial步骤二前往PyTorch官网获取安装命令这是最关键的一步永远以 PyTorch官网 的安装页面为准。页面上的选择器会帮你生成正确的命令。你的选择根据你的nvidia-smi结果。例如如果支持CUDA 12.4在官网选择PyTorch Build: Stable (2.3.0)Your OS: Windows/Linux/macOSPackage: Conda (或Pip但Conda方式通常更省心)Language: PythonCompute Platform: CUDA 12.1 (选择一个≤你驱动支持版本的稳定CUDA版本通常选推荐的最新版如12.1)选择后网站会给出类似这样的命令# Conda 方式 conda install pytorch torchvision torchaudio pytorch-cuda12.1 -c pytorch -c nvidia或者# Pip 方式 pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121步骤三执行安装与验证复制命令到激活的conda环境中执行。安装完成后启动Python解释器进行验证import torch print(torch.__version__) # 输出PyTorch版本如 2.3.0 print(torch.cuda.is_available()) # 输出 True 表示GPU可用 print(torch.cuda.get_device_name(0)) # 输出你的GPU型号如 ‘NVIDIA GeForce RTX 4090’如果torch.cuda.is_available()返回True恭喜你GPU环境配置成功如果返回False请按上述步骤检查驱动版本和安装命令的选择。注意对于使用Intel Arc显卡或Apple Silicon (M系列芯片) 的用户PyTorch也提供了相应的支持如DirectML后端或MPS后端但本教程以最主流的NVIDIA CUDA生态为例。这些用户请参考PyTorch官网针对特定硬件的安装指南。3. 张量Tensor一切计算的基石与核心操作PyTorch中所有数据的基本单位都是张量Tensor。你可以把它理解为NumPy的ndarray的升级版但关键区别在于PyTorch张量可以存储在GPU上加速计算并且内置了自动微分的跟踪能力。理解张量是理解一切的开始。3.1 张量的创建与基础属性张量可以有多种创建方式最常用的是直接从数据转换或使用工厂函数。import torch # 1. 从Python列表/NumPy数组创建 data [[1, 2], [3, 4]] x_np np.array(data) # 假设已导入numpy as np x_tensor_from_list torch.tensor(data) # 从列表创建 x_tensor_from_np torch.from_numpy(x_np) # 从NumPy数组创建共享内存 # 2. 使用工厂函数创建特定形状和内容的张量 zeros_tensor torch.zeros(2, 3) # 2行3列的全0张量 ones_tensor torch.ones(2, 3, 4) # 2x3x4的三维全1张量 rand_tensor torch.rand(2, 3) # 2x3的随机张量均匀分布0~1 randn_tensor torch.randn(2, 3) # 2x3的标准正态分布随机张量 # 3. 创建与现有张量属性一致的新张量 new_tensor torch.randn_like(randn_tensor) # 形状和dtype与randn_tensor一致 new_zeros torch.zeros_like(rand_tensor) # 形状一致的全0张量每个张量都有三个核心属性你需要时刻留意dtype: 数据类型如torch.float32,torch.int64,torch.bool。深度学习计算通常使用float32。shape: 张量的维度形状对应NumPy的.shape。device: 张量所在的设备cpu或cuda:0(GPU)。这是PyTorch与NumPy的核心区别之一。tensor torch.randn(3, 4, dtypetorch.float32) print(fShape: {tensor.shape}) # 输出: torch.Size([3, 4]) print(fData type: {tensor.dtype}) # 输出: torch.float32 print(fDevice: {tensor.device}) # 输出: cpu # 将张量移动到GPU如果可用 if torch.cuda.is_available(): tensor_gpu tensor.to(cuda) # 或者 tensor.cuda() print(fGPU Tensor device: {tensor_gpu.device}) # 输出: cuda:03.2 张量的操作索引、切片与变形张量的操作语法与NumPy高度相似学习成本极低。tensor torch.arange(12).reshape(3, 4) # 创建一个3x4的张量值为0-11 print(tensor) # 输出: # tensor([[ 0, 1, 2, 3], # [ 4, 5, 6, 7], # [ 8, 9, 10, 11]]) # 索引和切片 print(tensor[0]) # 第一行: tensor([0, 1, 2, 3]) print(tensor[:, 1]) # 第二列: tensor([1, 5, 9]) print(tensor[1:3, 2:]) # 第2-3行第3列到最后: tensor([[6, 7], [10, 11]]) # 连接张量 t1 torch.ones(2, 3) t2 torch.zeros(2, 3) t_cat_dim0 torch.cat([t1, t2], dim0) # 沿第0维行拼接形状变为 (4, 3) t_cat_dim1 torch.cat([t1, t2], dim1) # 沿第1维列拼接形状变为 (2, 6) # 改变形状 - view 和 reshape # view要求张量在内存中是连续的reshape更通用如果不连续会自动拷贝 tensor torch.arange(6) print(tensor.view(2, 3)) # 形状变为2x3 print(tensor.reshape(3, -1)) # -1表示自动计算该维度大小这里变为3x2 # 广播机制Broadcasting # 当两个张量形状不同时PyTorch会自动扩展较小张量的维度以进行计算 a torch.ones(3, 1, 2) # shape: (3, 1, 2) b torch.ones(1, 4, 2) # shape: (1, 4, 2) c a b # 广播后 shape: (3, 4, 2)每个元素都是 ab3.3 原地操作与非原地操作一个至关重要的细节这是PyTorch新手常踩的坑关系到内存管理和自动微分。非原地操作会返回一个新的张量原张量不变。大多数操作都是非原地的如add(),mul()。原地操作直接修改原张量的值操作符后带下划线_如add_(),mul_()。x torch.tensor([1., 2., 3.]) y torch.tensor([4., 5., 6.]) z x.add(y) # 非原地操作x不变z是新的张量 [5., 7., 9.] print(x) # tensor([1., 2., 3.]) x.add_(y) # 原地操作x被直接修改 print(x) # tensor([5., 7., 9.])实操心得在模型训练的前向传播中可以放心使用非原地操作。但在需要极致优化内存或性能的特定场景如大型模型推理原地操作能节省内存。但要注意过度使用原地操作可能会破坏计算图影响梯度计算在自定义网络层时需要谨慎。4. 自动微分AutogradPyTorch的“灵魂”引擎深度学习模型的训练本质是一个不断迭代的“预测-计算损失-反向传播梯度-更新参数”的过程。其中的“反向传播梯度”如果手动实现将是极其繁琐且容易出错的。PyTorch的autograd包自动完成了这个工作这是它最强大的特性之一。4.1 requires_grad与计算图的构建PyTorch通过跟踪在张量上执行的所有操作动态构建一个有向无环图DAG即计算图。图中的叶子节点通常是你的输入数据或模型参数中间节点是运算结果根节点是最终的输出如损失值。要让PyTorch跟踪某个张量的梯度你需要将其requires_grad属性设置为True。# 创建需要梯度的张量通常是模型参数 x torch.ones(2, 2, requires_gradTrue) # 方法1创建时指定 print(x.requires_grad) # True y torch.ones(2, 2) y.requires_grad_(True) # 方法2原地修改属性 print(y.requires_grad) # True # 一个简单的计算 z x 2 print(z) # tensor([[3., 3.], [3., 3.]], grad_fnAddBackward0)注意z的grad_fn属性它指向了创建z的加法函数AddBackward0。这表明z是计算图中的一个节点并且它知道如何计算关于其输入的梯度。4.2 backward()梯度计算与累积当我们完成前向计算得到一个标量损失必须是标量后调用.backward()方法autograd就会沿着计算图反向传播计算所有requires_gradTrue的张量叶子节点的梯度。x torch.tensor([1., 2., 3.], requires_gradTrue) y x 2 z y * y * 3 out z.mean() # 最终输出必须是一个标量这里是 z 的平均值 print(fOutput: {out}) # Output: tensor(63., grad_fnMeanBackward0) # 反向传播计算梯度 out.backward() # 梯度现在存储在 x.grad 中 print(fGradient of x: {x.grad}) # d(out)/dx # 计算过程out mean(3*(x2)^2) d(out)/dx 2*(x2)/len(x) * 3? 让我们手动验证 # 实际上对于 x_i: out 1/3 * Σ 3*(x_i2)^2 Σ (x_i2)^2 # d(out)/dx_i 2*(x_i2) # 当 x[1,2,3]时梯度应为 [2*(12), 2*(22), 2*(32)] [6, 8, 10] # 输出: tensor([6., 8., 10.])梯度累积这是一个关键点。在训练循环中我们通常在每次迭代batch的末尾调用optimizer.zero_grad()来清空梯度。如果不这样做下一次.backward()计算的梯度会与之前累积的梯度相加导致错误的参数更新。# 错误示例梯度累积 x torch.tensor(2.0, requires_gradTrue) for _ in range(3): y x ** 2 y.backward() # 每次反向传播梯度会累加到 x.grad print(fGrad after step: {x.grad}) # 输出可能类似: tensor(4.), tensor(8.), tensor(12.) 梯度在不断累加 # 正确做法每次迭代前清零梯度 x torch.tensor(2.0, requires_gradTrue) optimizer torch.optim.SGD([x], lr0.1) # 后续会讲优化器 for _ in range(3): optimizer.zero_grad() # 清零梯度 y x ** 2 y.backward() optimizer.step() # 根据梯度更新x print(fx after step: {x})4.3 梯度追踪的控制detach()与no_grad()在某些场景下我们不需要PyTorch跟踪梯度例如模型推理预测只需要前向计算不更新参数。评估指标计算如准确率不需要梯度。从计算图中提取中间值用于其他不需要梯度的计算。这时我们可以使用torch.no_grad()上下文管理器或.detach()方法。x torch.tensor([1., 2.], requires_gradTrue) y x * 2 # 方法1使用 .detach() 获得一个不需要梯度的新张量与原计算图分离 z_detached y.detach() print(z_detached.requires_grad) # False # 对 z_detached 的操作不会被跟踪 # 方法2使用 torch.no_grad() 上下文管理器块内所有计算都不会构建计算图 with torch.no_grad(): z_no_grad y * 3 print(z_no_grad.requires_grad) # False # 这在模型评估时非常有用 model ... # 你的神经网络 model.eval() # 将模型设置为评估模式影响Dropout、BatchNorm等层 with torch.no_grad(): for data, target in validation_loader: output model(data) # 计算准确率等指标...踩坑实录我曾经在计算一个自定义损失函数时不小心将requires_gradTrue的张量用于构造一个巨大的索引矩阵导致计算图异常庞大内存瞬间爆满。后来用.detach()将索引部分分离出来问题才解决。记住只有模型参数和需要优化目标直接依赖的输入才需要梯度。5. 神经网络构建用nn.Module搭建你的第一个模型PyTorch使用torch.nn模块来构建神经网络。其核心是nn.Module类。你自定义的任何一个网络层或整个网络都应该继承这个类。5.1 定义一个简单的全连接网络让我们构建一个用于手写数字识别MNIST数据集的简单网络。MNIST图像是28x28的灰度图我们将其展平为784维的向量作为输入经过两个全连接层输出10个类别的概率。import torch.nn as nn import torch.nn.functional as F class SimpleNN(nn.Module): def __init__(self, input_size784, hidden_size128, num_classes10): super(SimpleNN, self).__init__() # 必须调用父类初始化 # 定义网络层 self.fc1 nn.Linear(input_size, hidden_size) # 第一层全连接 self.fc2 nn.Linear(hidden_size, num_classes) # 第二层全连接 # 我们也可以定义Dropout等层 self.dropout nn.Dropout(p0.2) def forward(self, x): # 定义前向传播路径 # 输入x形状: (batch_size, 784) x x.view(-1, 784) # 确保输入被展平-1表示自动推断batch_size out self.fc1(x) # 形状: (batch_size, 128) out F.relu(out) # 使用ReLU激活函数注意这里用的是functional接口 out self.dropout(out) # 应用Dropout仅在训练时生效 out self.fc2(out) # 形状: (batch_size, 10) # 注意我们通常不在最后一层加激活函数因为损失函数如CrossEntropyLoss内部包含了Softmax return out # 实例化模型 model SimpleNN() print(model) # 输出网络结构 # SimpleNN( # (fc1): Linear(in_features784, out_features128, biasTrue) # (dropout): Dropout(p0.2, inplaceFalse) # (fc2): Linear(in_features128, out_features10, biasTrue) # )关键点解析__init__方法在这里定义网络中将要用到的所有“层”nn.Module的子类。nn.Linear,nn.Conv2d,nn.Dropout等都是预定义好的层。将它们赋值给selfPyTorch会自动注册这些层使得它们的参数可以被model.parameters()访问到从而被优化器更新。forward方法这里定义了数据从输入到输出的真实流动过程。你永远不需要直接调用forward方法而是通过调用模型实例output model(input)来执行前向传播。PyTorch会帮你调用forward并自动构建计算图。激活函数可以使用nn模块中的类如nn.ReLU()也可以使用torch.nn.functional中的函数如F.relu()。两者的区别在于nn.ReLU()是一个可以放在__init__中的层对象而F.relu()是一个纯函数通常在forward中直接调用。对于无参数的层如ReLU, Dropout两种方式都可以。5.2 访问模型参数与设备移动模型本身也是一个nn.Module我们可以方便地查看和操作它的参数。# 打印模型的所有参数 for name, param in model.named_parameters(): print(f{name}: {param.size()}) # 例如: fc1.weight: torch.Size([128, 784]) # 将模型移动到GPU device torch.device(cuda if torch.cuda.is_available() else cpu) model.to(device) # 这将递归地将所有模型参数和缓冲区移动到指定设备 # 之后输入数据也必须移动到相同设备 dummy_input torch.randn(32, 1, 28, 28).to(device) # 假设batch_size32 output model(dummy_input) # 前向传播在GPU上执行 print(output.device) # cuda:06. 损失函数与优化器驱动模型学习的两个轮子模型定义了从输入到输出的映射但如何让它输出正确的结果我们需要两个组件损失函数衡量输出与目标的差距优化器根据这个差距梯度来更新模型参数。6.1 常见的损失函数torch.nn模块提供了各种损失函数。nn.MSELoss均方误差损失常用于回归任务。nn.CrossEntropyLoss交叉熵损失这是分类任务最常用的损失函数。它内部集成了Softmax操作所以你的模型最后一层不需要再加Softmax激活。nn.BCELoss二分类交叉熵损失需要模型输出经过Sigmoid激活值在0-1之间。nn.BCEWithLogitsLoss结合了Sigmoid和BCELoss数值上更稳定推荐用于二分类。criterion nn.CrossEntropyLoss() # 对于十分类任务 # 假设一个batch的输出和标签 outputs model(dummy_input) # 形状: (32, 10)32个样本10个类别的得分logits labels torch.randint(0, 10, (32,)).to(device) # 模拟32个随机标签形状: (32,) # 计算损失 loss criterion(outputs, labels) print(fLoss: {loss.item()}) # .item() 将单元素张量转换为Python数字6.2 优化器从SGD到Adam优化器决定了如何利用梯度来更新参数。最基本的优化器是随机梯度下降SGD但实践中更常用的是其改进版如Adam。import torch.optim as optim # 实例化优化器将需要优化的参数model.parameters()和学习率传给它 optimizer optim.Adam(model.parameters(), lr0.001) # Adam优化器学习率0.001 # 或者使用SGD with Momentum # optimizer optim.SGD(model.parameters(), lr0.01, momentum0.9) # 在训练循环中优化器的标准三步曲 for epoch in range(num_epochs): for batch_idx, (data, target) in enumerate(train_loader): # 1. 清零梯度非常重要 optimizer.zero_grad() # 2. 前向传播 计算损失 data, target data.to(device), target.to(device) outputs model(data) loss criterion(outputs, target) # 3. 反向传播 loss.backward() # 4. 更新参数 optimizer.step() # 可选打印训练信息 if batch_idx % 100 0: print(fEpoch [{epoch1}/{num_epochs}], Step [{batch_idx1}/{len(train_loader)}], Loss: {loss.item():.4f})优化器选择心得SGD理论最基础调优好的SGD配合合适的学习率衰减和Momentum在最终性能上有时能超过Adam但它对学习率等超参数更敏感。Adam自适应学习率通常不需要精细调参就能获得不错的效果是快速实验和入门时的首选。对于大多数CV和NLP任务用Adam作为起点是个好选择。AdamWAdam的一个变种修正了权重衰减Weight Decay的实现在许多任务上尤其是Transformer类模型表现优于原始Adam是目前更推荐的选择。用法optim.AdamW(model.parameters(), lr0.001, weight_decay0.01)。7. 数据加载与预处理用DataLoader构建高效数据管道模型和优化器准备好了数据呢PyTorch提供了torch.utils.data.DataLoader和Dataset类来高效地加载和预处理数据。7.1 自定义Dataset你需要继承Dataset类并实现__len__和__getitem__方法。from torch.utils.data import Dataset, DataLoader from PIL import Image import os class CustomImageDataset(Dataset): def __init__(self, img_dir, transformNone): Args: img_dir (string): 图片目录路径。 transform (callable, optional): 一个可选的变换函数应用于样本。 self.img_dir img_dir self.img_names os.listdir(img_dir) # 假设目录下全是图片 self.transform transform def __len__(self): return len(self.img_names) def __getitem__(self, idx): img_path os.path.join(self.img_dir, self.img_names[idx]) image Image.open(img_path).convert(RGB) # 打开图片并转为RGB # 这里假设标签包含在文件名中例如 cat_001.jpg label self.img_names[idx].split(_)[0] # 将标签转换为整数索引这里需要你自己的标签映射逻辑 label_idx 0 if label cat else 1 # 简单示例猫为0狗为1 if self.transform: image self.transform(image) return image, label_idx7.2 使用Torchvision进行标准数据预处理对于常见数据集如MNIST, CIFAR-10, ImageNettorchvision库提供了现成的数据集和常用的图像变换Transforms。import torchvision.transforms as transforms import torchvision.datasets as datasets # 定义数据变换将图像转换为张量并归一化到[-1, 1]或[0, 1]区间 transform transforms.Compose([ transforms.ToTensor(), # 将PIL Image或numpy.ndarray转换为torch.Tensor并缩放到[0.0, 1.0] transforms.Normalize((0.5,), (0.5,)) # 对单通道灰度图像均值0.5标准差0.5使得值域在[-1,1] # 对于RGB三通道图像transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) # 加载MNIST数据集 train_dataset datasets.MNIST(root./data, trainTrue, transformtransform, downloadTrue) # 如果本地没有会自动下载 test_dataset datasets.MNIST(root./data, trainFalse, transformtransform)7.3 使用DataLoader进行批量加载DataLoader负责从Dataset中按批次抽取数据并提供了打乱shuffle、多进程加载num_workers等强大功能。# 参数详解 # batch_size: 每个批次的大小 # shuffle: 是否在每个epoch开始时打乱数据训练集通常为True测试集为False # num_workers: 用于数据加载的子进程数。在Linux/Mac上可以设置0加速Windows上有时设为0更稳定。 # pin_memory: 如果使用GPU设置为True可以将数据锁页内存加速GPU数据传输。 train_loader DataLoader(datasettrain_dataset, batch_size64, shuffleTrue, num_workers4, pin_memoryTrue) test_loader DataLoader(datasettest_dataset, batch_size64, shuffleFalse, # 测试集不需要打乱 num_workers4) # 使用示例 for batch_idx, (images, labels) in enumerate(train_loader): # images 形状: [64, 1, 28, 28] (batch_size, channels, height, width) # labels 形状: [64] # ... 将数据移动到设备喂给模型 ... if batch_idx 0: print(fBatch {batch_idx}: Images shape {images.shape}, Labels shape {labels.shape}) break性能调优提示num_workers的设置很关键。通常设置为CPU核心数或2-4倍。设置太高可能导致内存不足设置太低则数据加载会成为训练瓶颈。在Windows上由于多进程实现方式不同有时num_workers0反而更稳定。pin_memoryTrue在GPU训练时能带来明显的速度提升。8. 完整训练循环实战MNIST手写数字识别现在让我们把所有概念串联起来完成一个完整的、可运行的MNIST手写数字识别模型训练脚本。import torch import torch.nn as nn import torch.optim as optim import torchvision import torchvision.transforms as transforms from torch.utils.data import DataLoader # 1. 超参数定义 device torch.device(cuda if torch.cuda.is_available() else cpu) print(fUsing device: {device}) batch_size 64 learning_rate 0.001 num_epochs 5 input_size 784 # 28*28 hidden_size 128 num_classes 10 # 2. 数据加载 transform transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)) # MNIST数据集的均值和标准差 ]) train_dataset torchvision.datasets.MNIST(root./data, trainTrue, transformtransform, downloadTrue) test_dataset torchvision.datasets.MNIST(root./data, trainFalse, transformtransform) train_loader DataLoader(datasettrain_dataset, batch_sizebatch_size, shuffleTrue, num_workers4, pin_memoryTrue) test_loader DataLoader(datasettest_dataset, batch_sizebatch_size, shuffleFalse) # 3. 模型定义复用之前的SimpleNN class SimpleNN(nn.Module): def __init__(self, input_size, hidden_size, num_classes): super(SimpleNN, self).__init__() self.fc1 nn.Linear(input_size, hidden_size) self.relu nn.ReLU() self.fc2 nn.Linear(hidden_size, num_classes) # 注意没有在最后一层加Softmax def forward(self, x): x x.view(-1, input_size) # 展平图像 out self.fc1(x) out self.relu(out) out self.fc2(out) return out model SimpleNN(input_size, hidden_size, num_classes).to(device) # 4. 损失函数和优化器 criterion nn.CrossEntropyLoss() optimizer optim.Adam(model.parameters(), lrlearning_rate) # 5. 训练循环 total_step len(train_loader) for epoch in range(num_epochs): model.train() # 设置为训练模式影响Dropout、BatchNorm等层 running_loss 0.0 for i, (images, labels) in enumerate(train_loader): # 将数据移动到设备 images images.to(device, non_blockingTrue) labels labels.to(device, non_blockingTrue) # 前向传播 outputs model(images) loss criterion(outputs, labels) # 反向传播和优化 optimizer.zero_grad() loss.backward() optimizer.step() running_loss loss.item() if (i1) % 100 0: print(fEpoch [{epoch1}/{num_epochs}], Step [{i1}/{total_step}], Loss: {loss.item():.4f}) print(fEpoch [{epoch1}/{num_epochs}], Average Loss: {running_loss / total_step:.4f}) # 6. 在每个epoch后评估模型 model.eval() # 设置为评估模式 with torch.no_grad(): correct 0 total 0 for images, labels in test_loader: images images.to(device) labels labels.to(device) outputs model(images) _, predicted torch.max(outputs.data, 1) # 获取预测类别最大值的索引 total labels.size(0) correct (predicted labels).sum().item() print(fTest Accuracy after Epoch {epoch1}: {100 * correct / total:.2f} %) print(Training finished.) # 7. 保存模型 torch.save(model.state_dict(), mnist_simple_nn.pth) print(Model saved to mnist_simple_nn.pth)运行这个脚本你应该能看到损失逐渐下降测试准确率在5个epoch后能达到97%以上。这标志着你已经成功使用PyTorch完成了一个完整的深度学习项目流程。9. 模型保存、加载与推理训练好的模型需要保存下来以便后续部署或继续训练。9.1 保存与加载状态字典推荐PyTorch推荐只保存模型的state_dict()这是一个包含模型所有参数权重和偏置的字典。这种方式最灵活也最节省空间。# 保存 torch.save(model.state_dict(), model_weights.pth) # 加载 # 首先需要重新实例化模型结构必须和保存时完全一致 loaded_model SimpleNN(input_size, hidden_size, num_classes) # 然后将保存的参数加载到模型中 loaded_model.load_state_dict(torch.load(model_weights.pth)) loaded_model.to(device) loaded_model.eval() # 别忘了设置为评估模式9.2 保存整个模型不推荐也可以保存整个模型对象包括结构和参数但这种方式对代码结构有较强依赖且在不同PyTorch版本间可能不兼容。# 保存 torch.save(model, entire_model.pth) # 加载 loaded_model torch.load(entire_model.pth) loaded_model.to(device) loaded_model.eval()9.3 使用模型进行推理加载模型后进行推理预测的步骤很简单。# 假设我们有一张新的图片 new_image (PIL Image 或 numpy array) # 1. 应用与训练时相同的变换 transform transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)) ]) input_tensor transform(new_image).unsqueeze(0) # unsqueeze增加一个batch维度 - [1, 1, 28, 28] input_tensor input_tensor.to(device) # 2. 前向传播 with torch.no_grad(): output loaded_model(input_tensor) # output 形状: [1, 10] # 3. 获取预测结果 probabilities torch.nn.functional.softmax(output[0], dim0) # 转换为概率 predicted_class torch.argmax(probabilities).item() # 获取概率最大的类别索引 print(fPredicted digit: {predicted_class}) print(fProbabilities: {probabilities.cpu().numpy()}) # 转换为numpy数组查看10. 调试技巧与常见陷阱即使理解了所有概念实际编码中依然会遇到各种问题。这里分享几个我踩过坑后总结的调试技巧。1. 张量形状不匹配RuntimeError: shape mismatch这是最常见的错误。养成打印张量形状的习惯在模型forward函数的关键步骤、损失计算前插入print(x.shape)。确保每一层的输入输出维度符合你的预期。例如全连接层nn.Linear(in_features, out_features)要求输入的最后维度是in_features。2. 忘记调用zero_grad()这会导致梯度累积使训练完全失控。一个良好的习惯是在loss.backward()之前立即写上optimizer.zero_grad()。3. 在训练和评估模式间切换model.train()和model.eval()不仅仅是一个标识。它们会改变Dropout和BatchNorm等层的行为。在训练循环开始时调用model.train()在验证和测试时调用model.eval()并用with torch.no_grad():包裹评估代码。4. 设备不一致RuntimeError: Expected all tensors to be on the same device确保模型、输入数据、标签都在同一个设备上CPU或GPU。一个可靠的模式是device torch.device(cuda if torch.cuda.is_available() else cpu) model Model().to(device) # 在数据加载循环中 for data, target in dataloader: data, target data.to(device), target.to(device) # ...5. 损失不下降或为NaN学习率太大尝试降低学习率如从0.001调到0.0001。数据未归一化输入数据值域过大如0-255的像素值会导致梯度爆炸。务必使用transforms.Normalize。网络结构问题检查激活函数是否合理如分类最后一层不要用ReLU梯度是否消失对于深网络考虑使用ResNet结构或更好的初始化。损失函数选择错误确保损失函数与任务匹配如分类用交叉熵回归用MSE。6. 使用torch.no_grad()节省内存和计算在不需要计算梯度的场景如模型评估、特征提取务必使用with torch.no_grad():。这会禁用梯度计算和存储大幅减少内存占用并提升速度。掌握这些核心概念和实操步骤你已经具备了使用PyTorch进行独立深度学习项目开发的基础能力。接下来的方向可以是深入计算机视觉CNN、自然语言处理RNN, Transformer或是探索更高级的自动微分、自定义算子等功能。PyTorch的生态非常丰富但万变不离其宗理解本文所述的张量、自动微分、模块和数据管道将是你驾驭这一切的基石。