十年匠心定制 · 商业建站与技术教学双线并行 咨询热线:400-886-1026 service@lmnt.cn
ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

张量计算陷阱全解析:从NaN、显存泄漏到设备不一致

张量计算陷阱全解析:从NaN、显存泄漏到设备不一致 训练深度学习模型时最让人头疼的往往不是网络结构设计不出来而是训练到一半 loss 变成nan或者换了一台机器结果完全对不上。这类问题看起来像是“玄学”实际上绝大多数都出在张量计算这一层数据精度不对、设备不统一、梯度异常、显存没有释放甚至只是少写了一行zero_grad()。tensor_of_ice这个名字正好概括了这类问题的本质——张量层面的“冻结”数值冻结、显存冻结、模型状态冻结。任何一个环节被“冻住”整个训练流程就停摆了。这篇文章不是某个神秘开源仓库的说明书而是站在工程落地的角度把张量计算里最容易被忽略的细节拆开讲透。你会看到张量的基础概念、环境搭建、完整训练示例、数值稳定性问题、常见报错排查方法以及在生产环境中常见的工程建议。读完这篇文章你应该能回答这几个问题为什么我的 loss 突然变成 NaN为什么 GPU 显存越跑越大为什么同一个模型在不同机器上结果不一样以及怎样从一开始就避免这些坑。1. 这篇文章真正要解决的问题很多初学者会把“学会 PyTorch”等同于“会写nn.Module”于是很快陷入一个误区模型结构抄过来了损失函数也写对了但是训练出来的结果完全不可用。实际上现代深度学习框架已经把网络搭建的复杂度降到很低真正的复杂度在张量计算和工程细节上。具体来说tensor_of_ice想解决的是下面这三类问题。第一类是数值冻结。训练过程中 loss 突然变成nan或者inf梯度不再更新模型参数像被冻住一样。这类问题通常与精度下溢、上溢、除零、日志函数输入不合法有关。第二类是显存冻结。显存占用随训练轮数持续上升最终 OOM。这背后是计算图中保存了不需要的中间变量或者张量没有及时释放。第三类是模型状态冻结。模型在保存和加载之后表现不一致或者分布式场景下设备不对齐导致模型“看起来在跑实际已经坏了”。因此这篇文章适合这样的读者刚开始用 PyTorch 做深度学习想弄明白tensor、requires_grad、device这些基础概念已经能跑通简单模型但遇到 NaN、OOM、结果不可复现等问题想建立一套系统的排查思路准备把训练代码部署到生产环境需要了解张量层面的最佳实践和性能优化手段。一句话总结本文的核心判断模型结构决定效果上限张量计算决定工程下限。如果你的张量层不稳定再好的模型结构也落不了地。2. 张量的基础概念与核心原理2.1 什么是张量张量Tensor是深度学习框架中最基础的数据结构。从数学角度看它是标量、向量、矩阵在高维空间的推广标量Scalar是 0 维张量向量Vector是 1 维张量矩阵Matrix是 2 维张量3 维及以上的数组统称为 N 维张量。从工程角度看张量就是一个多维数组但它比普通数组多了几个关键属性shape形状、dtype数据类型、device所在设备、requires_grad是否需要梯度。正是这些属性让张量能在 GPU 上高效执行矩阵运算并且支撑自动求导。2.2 张量与 NumPy 数组的区别很多初学者会问既然 NumPy 也有 ndarray为什么还要用 Tensor关键区别在于两点。一是自动求导。Tensor 通过requires_gradTrue记录运算过程前向传播时保存计算图反向传播时自动计算梯度。这是模型训练的基石。二是硬件加速。Tensor 可以无缝地在 CPU 和 GPU 之间迁移配合 CUDA 获得数十倍的矩阵运算加速。NumPy 数组则主要运行在 CPU 上。对比维度NumPy ndarrayPyTorch Tensor多维数组支持支持GPU 加速不支持需额外库原生支持自动求导不支持支持神经网络接口无nn.Module等高层接口设备管理无to(device)显式管理2.3 贯穿全文的“冰”的隐喻用“冰”来理解张量计算中的三类不稳定问题非常直观。数值之冰一个极小的数在连续相乘后可能下溢为 0一个极大的数在指数运算中可能上溢为inf。这些都像水结成冰一样让计算过程突然失去流动性。解决的关键是稳定数值计算例如 softmax 减去最大值、log 域计算等。显存之冰训练过程中如果没有及时释放中间张量显存就一点点被“冻住”最终 OOM。这个问题的本质是计算图中的引用关系没有被正确管理。状态之冰模型保存、加载、设备迁移、随机种子设置不当会让模型处于“状态不一致”的冻结状态。同一个模型昨天跑和今天跑结果可能完全不同。理解这三层隐喻后续的代码示例和排查思路就有了统一的线索。3. 环境准备与前置条件在开始写代码之前先把运行环境准备好。本文的示例以 PyTorch 为主因为它在学术和工业界的使用都很广泛API 设计也更贴近工程实践。如果你使用 TensorFlow核心思路仍然适用只是 API 名称不同。3.1 安装 Python 虚拟环境建议使用venv或conda创建独立环境避免不同项目的依赖互相污染。以venv为例python -m venv ./venv source ./venv/bin/activate # Linux/macOS # .\venv\Scripts\activate # WindowsPython 版本建议 3.8 及以上。具体版本请以实际项目依赖为准本文重点是通用思路。3.2 安装 PyTorchPyTorch 的安装命令会根据你的操作系统、包管理工具和 CUDA 版本而变化。最稳妥的方式是去 PyTorch 官网选择对应配置生成安装命令。一个不含 CUDA 的 CPU 版本示例pip install torch torchvision torchaudio如果你有 NVIDIA GPU则需要在官网选择 CUDA 版本。安装完成后用下面的代码验证环境是否可用。# 文件路径env_check.py import torch print(PyTorch 版本:, torch.__version__) print(CUDA 是否可用:, torch.cuda.is_available()) print(GPU 数量:, torch.cuda.device_count()) if torch.cuda.is_available(): print(GPU 名称:, torch.cuda.get_device_name(0))运行命令python env_check.py如果torch.cuda.is_available()返回False并不代表环境不能工作只是当前环境没有可用的 CUDA 设备。后续示例在 CPU 上同样可以运行只是训练速度会慢一些。3.3 固定随机种子为了让实验结果可复现训练开始时建议固定随机种子。否则每次运行生成的随机数不同模型结果会有差异。这个问题在调参和代码评审时经常被忽略。import random import numpy as np import torch def set_seed(seed: int 42): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)4. 核心流程拆解训练一个深度学习模型本质上是围绕张量做五件事准备数据、构建计算图、前向传播、反向传播、更新参数。每一步都有对应的张量操作也都有各自的坑。4.1 数据转换为张量原始数据通常是文件、数据库记录或 NumPy 数组并不能直接进入模型。第一步是把数据转换成张量。这一步最容易犯的错误是数据类型不匹配。例如 PyTorch 默认的浮点类型是torch.float32但某些模型导出或数据集转换时数据可能变成了torch.float64。在 CPU 上运行可能没有问题一旦切换到 GPU部分操作会报类型错误。import torch import numpy as np arr np.array([1.0, 2.0, 3.0]) tensor_from_numpy torch.from_numpy(arr) print(tensor_from_numpy.dtype) # torch.float64 # 转换为模型常用精度 tensor_float32 tensor_from_numpy.float()4.2 构建计算图与 requires_grad张量最强大的能力是自动求导。当你把一个张量的requires_grad设为TruePyTorch 会记录所有对该张量的操作形成一个计算图。反向传播时梯度会沿着计算图自动回传。这个机制的坑在于不是所有的张量都需要梯度。例如输入数据、标签、不需要学习的中间变量如果错误地开启了requires_grad会无谓地扩大计算图增加显存和耗时。正确做法是只在模型参数上开启梯度这通常由nn.Module自动完成。4.3 前向传播与损失计算前向传播是把输入张量经过各层变换得到预测结果。随后计算预测值与真实标签之间的损失。这里常见的坑是形状不对齐。例如一个(batch, seq_len)的张量和一个(batch,)的张量直接相减在部分维度上会触发广播机制但语义可能是错的。建议在关键位置打印张量的shape确认对齐后再继续。4.4 反向传播与前梯度的清空loss.backward()会根据计算图计算梯度。但 PyTorch 默认是累加梯度而不是替换梯度。这意味着如果你在同一个模型上多次调用backward()梯度会叠加。所以每一轮训练开始前必须调用optimizer.zero_grad()清空上一轮梯度。很多初学者的 NaN 问题就是忘记清空梯度导致的梯度不断累加数值越来越大最终溢出为inf或nan。4.5 参数更新optimizer.step()会根据梯度更新模型参数。这一步本身不复杂但要注意它和zero_grad()、backward()的调用顺序。标准顺序是optimizer.zero_grad()loss.backward()optimizer.step()顺序错乱的典型表现是loss 下降缓慢或者参数不更新。5. 完整示例与代码实现这一部分用 6 个示例串起从张量基础到模型训练的完整链路。每个示例都可以独立运行也建议按顺序执行便于理解。5.1 示例 1环境检查这个示例用来检查 PyTorch 和 CUDA 环境是否正常。# 文件路径env_check.py import torch print(PyTorch 版本:, torch.__version__) print(CUDA 是否可用:, torch.cuda.is_available()) print(GPU 数量:, torch.cuda.device_count()) if torch.cuda.is_available(): print(GPU 名称:, torch.cuda.get_device_name(0))运行python env_check.py如果 CUDA 不可用说明当前环境没有 GPU 或 PyTorch 未正确安装 CUDA 版本。后续示例仍可在 CPU 上运行。5.2 示例 2张量运算与自动求导这是理解requires_grad和计算图的最小示例。# 文件路径autograd_demo.py import torch x torch.tensor([2.0, 3.0], requires_gradTrue) y x * x z y.sum() z.backward() print(x , x) print(z , z.item()) print(x.grad , x.grad) # dz/dx 2*x预期输出x tensor([2., 3.], requires_gradTrue) z 13.0 x.grad tensor([4., 6.])关键逻辑z x1^2 x2^2所以dz/dx1 2*x1 4dz/dx2 2*x2 6。backward()执行后梯度被写入了x.grad。5.3 示例 3线性回归完整训练用线性回归演示完整的“数据准备 → 模型构建 → 训练循环 → 参数更新”流程。# 文件路径linear_regression.py import torch import torch.nn as nn import torch.optim as optim torch.manual_seed(42) # 1. 构造模拟数据 X torch.randn(1000, 3) w_true torch.tensor([2.0, -3.0, 5.0]) y X w_true 1.0 0.1 * torch.randn(1000) # 2. 定义模型 model nn.Linear(3, 1) loss_fn nn.MSELoss() optimizer optim.SGD(model.parameters(), lr0.01) # 3. 训练循环 for epoch in range(50): optimizer.zero_grad() pred model(X) loss loss_fn(pred.squeeze(), y) loss.backward() optimizer.step() if epoch % 10 0: print(fepoch {epoch}, loss{loss.item():.4f}) # 4. 输出学习到的参数 print(学习到的权重:, model.weight.data.squeeze()) print(学习到的偏置:, model.bias.data)预期输出数值存在微小波动epoch 0, loss... epoch 10, loss... ... 学习到的权重: tensor([ 1.98, -2.98, 4.96]) 学习到的偏置: tensor([1.01])这个示例的核心价值在于它完整展示了张量在训练循环中的生命周期。model(X)在前向传播时构建计算图loss.backward()沿计算图回传梯度optimizer.step()更新参数下一轮再从zero_grad()开始。5.4 示例 4数值稳定性问题与优化前面提到的“数值之冰”最常见的场景是 softmax 计算。初学者会写出下面这种“教科书版”实现# 文件路径softmax_demo.py import torch def unstable_softmax(logits): exp torch.exp(logits) return exp / exp.sum(dim-1, keepdimTrue) logits torch.tensor([1000.0, 1001.0, 1002.0]) print(unstable_softmax(logits))运行后torch.exp(1000)会溢出为inf最终结果变成[nan, nan, nan]。这就是典型的数值冻结问题。稳定做法是使用LogSumExp技巧先减去最大值再计算指数这样所有项都落在安全的数值范围内。def stable_softmax(logits): shifted logits - logits.max(dim-1, keepdimTrue).values exp torch.exp(shifted) return exp / exp.sum(dim-1, keepdimTrue) logits torch.tensor([1000.0, 1001.0, 1002.0]) print(stable_softmax(logits))预期输出tensor([0.0900, 0.2447, 0.6652])你可能会想既然 PyTorch 内置了F.softmax为什么还要理解这个技巧因为类似问题会出现在注意力机制、损失函数、自定义算子等地方。理解数值稳定性的原理能帮你在遇到 NaN 时快速定位问题而不是干着急。另一个常见稳定性手段是梯度裁剪将梯度范数限制在可控范围内# 在 backward 之后、step 之前调用 torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0)5.5 示例 5设备迁移与显存清理深度学习项目中设备不一致是高频问题。一个典型的报错是Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!根因是模型在 GPU 上输入数据却还在 CPU。解决方法是把模型和输入都迁移到同一个设备。# 文件路径device_demo.py import torch device torch.device(cuda if torch.cuda.is_available() else cpu) print(当前使用设备:, device) X torch.randn(16, 3).to(device) model torch.nn.Linear(3, 1).to(device) y model(X) print(输出张量所在设备:, y.device) # 显存清理GPU 场景 import gc gc.collect() torch.cuda.empty_cache()需要提醒的是torch.cuda.empty_cache()只是释放 PyTorch 的缓存显存并不会立刻把显存占用降为 0也不会解决未释放的引用问题。真正重要的是避免在计算图中保留不必要的中间张量。5.6 示例 6模型保存、加载与评估模型的保存和加载是“状态之冰”的典型场景。正确做法是只保存模型参数而不是整个模型对象。只保存参数的方式兼容性更好即使模型类定义发生小幅修改也能尽量恢复。# 文件路径save_load_demo.py import torch import torch.nn as nn model nn.Linear(3, 1) # 保存模型参数 torch.save(model.state_dict(), model.pt) # 在需要时加载模型参数 new_model nn.Linear(3, 1) new_model.load_state_dict(torch.load(model.pt, weights_onlyTrue)) new_model.eval() print(模型参数加载完成)weights_onlyTrue是 PyTorch 新版本的推荐写法可以提高加载安全性。加载后调用eval()会根据需要切换模型中的 dropout 和 batch normalization 行为确保推理结果稳定。6. 运行结果与效果验证6.1 如何判断训练是否成功以示例 3 的线性回归为例判断成功有几个标准loss 持续下降每一轮的 loss 总体呈下降趋势而不是大幅波动或保持不变。参数逼近真实值学习到的权重接近构造数据时使用的w_true [2.0, -3.0, 5.0]偏置接近1.0。无 NaN/Inf 出现loss 始终保持有限数值。如果 loss 不降反升先检查学习率是否过大如果 loss 变成nan优先检查数据中是否有 NaN或者梯度是否溢出。6.2 可复现性验证深度学习实验需要可复现否则团队协作会非常痛苦。在训练脚本开头固定随机种子后同一份代码在相同环境下两次运行的 loss 曲线应当完全一致。如果你的结果仍然不一致检查是否使用了 DataLoader 的随机 shuffling以及是否引入了非确定性的 GPU 算子。6.3 失败时第一步看哪里训练失败时不要直接怀疑模型结构。推荐的排查顺序是打印输入数据的shape和dtype确认数据是否正确进入模型打印 loss 是否有限值如果不有限尝试用更简单的数据或更小的学习率打印每个参数的grad是否为None或nan检查张量是否都在同一个设备上检查显存占用是否持续上涨。这个顺序能覆盖 80% 以上的“玄学”问题。7. 常见问题与排查思路以下是张量计算中最常见的 8 类问题按出现频率排序。问题现象可能原因排查方式解决方案loss 变成 NaN/Inf数值溢出、数据含 NaN、学习率过大检查输入数据、打印梯度、尝试降低学习率数据预处理、梯度裁剪、使用稳定数值算法梯度为 None参数没有参与计算图检查参数是否被detach或在无梯度上下文调整计算逻辑确保参数参与前向传播设备不一致报错模型与输入张量不在同一设备打印.device属性统一调用.to(device)运行时 shape 不匹配前向传播中维度对齐错误打印每层输入的 shape对照模型输入输出维度修正张量变换显存占用持续上涨计算图保留了中间张量监控显存曲线使用with torch.no_grad()推理、清理循环内无用张量模型加载后结果不对未调用eval()或参数加载不完整打印 state_dict 键名确认加载状态加载后调用model.eval()两次运行结果不一致随机种子未固定对比日志中的随机数固定 seed限制 DataLoader 随机性训练速度过慢GPU 未生效或 batch size 过小检查torch.cuda.is_available()合理利用.to(device)适当增大 batch size需要说明的是显存占用持续上涨不一定都是代码 bug还有一种常见原因是在训练循环里使用了loss.item()之外的张量切片导致 Python 层面持有张量引用无法被垃圾回收。排查时可以把循环内的张量赋值改用 Python 标量。8. 最佳实践与工程建议8.1 统一 dtype 与 device 规范在项目一开始就定下规范模型参数、输入数据统一使用torch.float32在训练入口统一完成.to(device)。不要在每个函数里临时做设备转换那样很容易遗漏导致设备不一致报错。对于大规模训练再考虑混合精度训练让部分计算使用float16以提升吞吐。8.2 固定随机种子与实验记录训练脚本中加入set_seed函数并记录每一次实验的配置、代码版本、数据版本和随机种子。这在团队协作和模型复现中非常重要。哪怕只是个人研究固定 seed 也能让你在调整超参数时确认结果差异确实来自超参数而不是随机波动。8.3 及时保存 checkpoint训练过程中定期保存模型参数到磁盘建议保留最近 2-3 个版本。训练中断时可以从最近的 checkpoint 恢复而不是从头开始。同时保留最优模型和最后一轮模型便于对比。# 文件路径checkpoint_demo.py import torch # 保存 torch.save({ epoch: epoch, model_state_dict: model.state_dict(), optimizer_state_dict: optimizer.state_dict(), loss: loss, }, checkpoint.pt) # 加载 checkpoint torch.load(checkpoint.pt, weights_onlyTrue) model.load_state_dict(checkpoint[model_state_dict]) optimizer.load_state_dict(checkpoint[optimizer_state_dict])8.4 训练循环之外用 no_grad 提高效率模型推理、验证集评估等不需要梯度的场景应该使用torch.no_grad()包裹避免构建不必要的计算图既节省显存也提升速度。with torch.no_grad(): val_pred model(val_X) val_loss loss_fn(val_pred.squeeze(), val_y)8.5 梯度裁剪与安全边界对 RNN、Transformer 等模型梯度爆炸是常见问题。在backward()之后step()之前加入梯度裁剪能显著提升训练稳定性。同时要关注数据隐私和模型安全问题训练数据需经过合法授权模型文件不应包含未脱敏的个人信息。8.6 从最小规模开始验证一开始不要用全量数据训练。先用 100 条数据、单 batch 跑通确认前向传播、反向传播、参数更新都没有问题后再逐步扩展数据规模。这个习惯能节省大量排查时间也是生产环境变更前必须先在测试环境验证的原则。9. 总结与后续学习方向张量计算是深度学习中“看不见的地基”。模型结构再好如果张量的数值、精度、设备、显存管理出了问题项目一样会卡壳。本文从tensor_of_ice这个隐喻切入说明张量层面出现“冻结”时会表现为数值 NaN、显存上涨、状态不一致也通过完整代码示例演示了从环境准备、张量运算、模型训练到保存加载的完整链路。下一步如果你希望继续深入可以沿着这几个方向展开一是深入理解autograd的实现机制弄清楚计算图到底怎么存储和释放二是学习混合精度训练与分布式训练中的张量通信三是研究模型量化与张量部署优化把张量层面的能力转化为真实业务价值。在所有动手实践之前先记住一句话先在最小规模上跑通再追求规模先把张量层的稳定性管住再谈模型效果。这条原则能帮你避开大部分深度学习工程中的暗礁。建议把本文收藏备用遇到张量相关的问题时按照文中第 7 章的排查表依次检查会比盲目改模型结构高效得多。
返回列表