
调试 PyTorch 模型的时候很多人只盯着 loss 曲线和 accuracy模型一旦训练完成但结果不对就不知道该从哪里查。这时最值得做的一件事就是把神经网络中间层输出可视化出来。所谓中间层输出就是输入经过某一层之后生成的张量对 CNN 来说通常是卷积层的 feature map对注意力模型来说可能是 attention weight对 RNN 来说就是隐藏状态。可视化之后你不再只能靠日志里的数字猜原因而是直接看图判断模型在哪一层提取了什么特征哪一层开始失效输入预处理是不是一开始就错了。我会按实际调试顺序拆一遍环境准备、forward_hook 捕获中间层输出、特征图保存、卷积核可视化、TensorBoard 记录最后是常见问题排查。这套方法对刚入门 PyTorch 的读者以及训练里遇到看不清问题的开发者都有用。你不需要改模型结构只靠一个 hook 和几行保存代码就能跑起来。1. 先搞清楚为什么要可视化神经网络中间层输出1.1 中间层输出到底指什么要把中间层输出可视化讲清楚先要理解一个事实神经网络不是一个只能看输入输出的黑盒。模型内部每一层都会生成中间张量。以 CNN 为例卷积层输出的 feature map 形状是 [batch, channels, height, width]每个通道相当于一个特征响应图数值大的地方表示该卷积核在该区域被激活了。激活函数层的输出、池化层的输出、注意力层的 attention weight、RNN 的 hidden state 也都算中间层输出。很多模型调试场景里单独看最终输出是看不出来的。比如分类错误可能有很多原因输入预处理不对训练数据分布和预处理不一致模型某层在梯度更新后失效或者自定义模块里的 shape 传递错了。中间层输出能把这些问题拆开第一层 feature map 乱大概率是输入问题中间层 feature map 开始出现大片噪声可能是训练异常或结构问题最后一层 feature map 没有区分度则说明特征提取或者分类头出了问题。1.2 可视化能解决哪些实际问题我一般会在下面几种场景里使用中间层可视化输入预处理检查。想确认 ImageNet 预训练模型的 mean/std 是否设置正确第一步看第一个卷积层的 feature map。自定义模型调试。自己拼的模块经常出现通道数对不上、输出 shape 和预期不同的问题可视化能快速确认每一层是不是真的输出了预期大小的张量。训练异常分析。loss 下降但是准确率不涨或者 loss 直接变成 nan这时候看中间层的数值分布比盲改学习率更有方向。模型结构理解。学习阶段看 ResNet 各阶段 feature map能明显看出不同层提取的是边缘、纹理还是语义信息。注意可视化中间层输出并不等于自动解释模型。它只是给你一套“看模型内部状态”的手段。真正判断问题还要结合输入样例、loss 曲线和网络结构一起分析。1.3 并不是所有层都需要可视化不要全网给每一个模块都加 hook。像 ResNet18 这种结构如果对所有 Conv2d 都注册 hook一次 forward 会捕获大量中间输出内存和磁盘压力都不小而且太多图像反而看不出重点。我的习惯是第一个卷积层看输入是否正常中间某一两个代表性 block 看特征提取是否正常最后一个卷积层或分类层之前的输出看特征是否有语义区分。先看 2 到 3 个层确认全链路没问题之后再决定要不要看更多层。建议第一次跑通时只选一个浅层和一个深层。浅层用来看输入和边缘特征深层用来看模型最终的语义特征。2. 环境准备用最小依赖跑通第一次特征图保存2.1 PyTorch 环境搭建的通用判断可视化中间层输出本身不需要很重的环境。核心依赖只有三个PyTorch、torchvision、matplotlib。如果你要记录训练过程中的特征图可以再加一个 torch.utils.tensorboard它是 PyTorch 自带的扩展装了 tensorboard 就能用。如果你还没有 PyTorch 环境建议先用 conda 或 venv 单独建一个环境不要和系统 Python 混在一起conda create -n vis_demo python3.10 conda activate vis_demo pip install torch torchvision matplotlib tensorboard这个安装方式在没有 GPU 的机器上也能用。CPU 版本跑一次 ResNet18 前向只需要几百毫秒到几秒具体取决于图片大小和机器性能用于可视化调试完全足够。如果你的机器有 NVIDIA GPU先执行nvidia-smi看驱动和 CUDA 版本再按官网说明安装对应 GPU 版 PyTorch。注意不是所有环境都能直接 pip 装到 GPU 版离线环境要格外注意版本匹配这点后面排查部分会提到。2.2 准备一张测试图片和输入预处理可视化需要一张真实图片不要用随机噪声。随机噪声也能生成 feature map但看不出语义不太适合调试输入预处理。准备方式有两种一种是从本地选一张普通图片另一种是用代码临时生成一张有内容的结构化色块图。我建议先用自己的图片结构越清晰越好比如包含明显的物体边缘、文字或几何图形。图片准备好之后要做标准化预处理。以 ResNet18 为例它在 ImageNet 上训练所以推荐使用下面的 transformfrom PIL import Image 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]) ]) image Image.open(demo.jpg).convert(RGB) input_tensor transform(image).unsqueeze(0) # [1, 3, 224, 224]这里有两个容易出错的地方。第一Resize 之后要确认图片被转换成 224x224否则后面打印 shape 时会和预期不一致。第二Normalize 操作必须在 ToTensor 之后因为 ToTensor 把 PIL 图片从 HWC 转成 CHW并且把像素值缩到 [0,1]之后 Normalize 才能用 ImageNet 的均值和标准差。2.3 加载预训练模型并设置 eval 模式中间层可视化多数时候只是推理不是训练。所以加载模型后必须调用model.eval()原因有两个Dropout 在 train 模式下会随机丢弃节点导致两次 forward 同一张图结果不一样。BatchNorm 在 train 模式使用当前 batch 的统计量在 eval 模式使用训练期间累计的全局统计量。对于预训练模型不切 eval 的话推理结果可能不稳定。加载方式我一般写成兼容模式import torch import torchvision.models as models model models.resnet18(pretrainedTrue) model.eval()如果你的 torchvision 版本比较新旧写法可能会提示空权重名称警告可以改成model models.resnet18(weightsmodels.ResNet18_Weights.DEFAULT)第一次执行时会联网下载预训练权重文件不大但网络不好的环境会比较慢。如果只是想验证 hook 流程也可以不加载预训练权重直接用随机初始化的模型model models.resnet18(weightsNone)随机初始化的 feature map 没有语义但能确认代码流程和 shape 是否正确。这个办法很适合先排除“环境、依赖、代码写错”的问题再切换回预训练权重看真实效果。3. 捕获中间层输出forward_hook 机制详解3.1 为什么推荐用 hook 而不是改 forward很多人遇到“想拿中间层输出”的第一个方案是把模型源码复制出来改 forward在里面加 print 或 return。这个做法在学习和试验阶段能跑但并不好用一是要维护一份模型副本和原模型结构容易脱节二是改 forward 会影响模型的正常调用尤其是模型在别处已经被封装过时三是一旦换了项目模型又要重新复制一套。PyTorch 提供了更干净的方式register_forward_hook也就是前向钩子。注册后模型的前向计算不会被改动但每次该模块执行完 forward 时都会调用你注册的回调函数你可以在回调里拿到输入和输出张量。这个机制非常适合做中间层可视化。3.2 最小代码注册 forward_hook 并触发以 ResNet18 为例我想看 layer1 和 layer4 的输出。先准备一个存储字典和 hook 工厂函数feature_map {} def make_hook(name): def hook(module, input, output): feature_map[name] output.detach() return hook model.layer1.register_forward_hook(make_hook(layer1)) model.layer4.register_forward_hook(make_hook(layer4)) with torch.no_grad(): final_output model(input_tensor) print(feature_map[layer1].shape) # [1, 64, 56, 56] print(feature_map[layer4].shape) # [1, 512, 7, 7]解释一下 hook 函数的三个参数module触发钩子的模块对象比如model.layer1。input该模块的输入。多数情况下是 tuple因为模块可能接受多个参数这里不打算改输入所以不处理。output该模块的返回值。对 Conv2d 和 Sequential 来说通常是 Tensor可以直接拿到特征图。不要忘记加with torch.no_grad()。只是可视化特征图时我们不需要保存梯度关掉梯度既减少显存占用也避免因为保存大量中间计算图导致后续 OOM。3.3 关键参数和数据的生命周期上面的 hook 里我调用了output.detach()这步很关键。如果不 detachoutput 会保留在计算图中虽然torch.no_grad()下前向不会构建计算图但如果你在训练循环中可视化同一个层并且没有 detach内存会越来越大。把 detach 这一步看成“只保留当前值不保留计算关系”。feature_map是一个普通字典。每次模型 forwardhook 都会执行一次把该层的 output 覆盖到字典里。如果你跑多个 batch字典里永远只保留最后一次 forward 的结果。这是故意的因为我们一般只看某一次输入或某个 step 的特征。如果你想把多张图的输出都保留可以改成列表按 index 或 step 追加但要控制数量。注册 hook 时会返回一个 handle。在训练脚本里如果反复注册同一个模块的 hook会出现 hook 累积。比如你在一个 for 循环里注册 100 次一次 forward 会触发 100 次回调。长期运行或做性能测试时建议保留 handle在不需要时调用 remove()handle model.layer1.register_forward_hook(make_hook(layer1)) # 不需要时 handle.remove()3.4 注册多个层时要注意什么用 named_modules 可以一次性遍历所有模块按类型筛选handles [] for name, module in model.named_modules(): if isinstance(module, torch.nn.Conv2d): handles.append(module.register_forward_hook(make_hook(name)))这个写法方便但要注意卷积层的数量很多一次 forward 会存下很多张特征图内存占用明显上升而且后面看图像时会看不过来。更有针对性一点可以先打印模型结构只看自己关心的层for name, module in model.named_modules(): print(name, type(module).__name__)然后选 2 到 3 个关键层注册 hook。不要一上来就全网打点。批量注册更适合你确定要看全部层的输出趋势且显存和磁盘都够的情况下。4. 特征图可视化从 tensor 到图片4.1 张量转换的固定三步拿到 feature map 之后它还是一个 PyTorch Tensor。要画成图片需要执行固定三步tensor feature_map[layer1] # [1, 64, 56, 56] feat tensor[0].detach().cpu().numpy() # [64, 56, 56]如果模型在 GPU 上必须先.cpu()再.numpy()。如果模型在 CPU 上.cpu()也不会报错所以干脆固定写成detach().cpu().numpy()避免 GPU 环境踩坑。layer1 输出的 shape 是 [batch, channel, height, width]取 batch 维度之后就是 [channel, height, width]。layer4 输出 [1, 512, 7, 7]取出来是 [512, 7, 7]。通道数很多分辨率很小所以要按下面的方式做归一化和拼接。4.2 灰度归一化的正确写法CNN 的 feature map 数值不是天然的像素值。它可能是负数也可能超过 1直接丢给 imshow 会得到一团黑或者一团白。需要用归一化把值域映射到 [0, 1] 之间。最简单的做法是逐通道 min-max 归一化import numpy as np def normalize_channel(ch): ch_min ch.min() ch_max ch.max() if ch_max - ch_min 1e-8: return np.zeros_like(ch) return (ch - ch_min) / (ch_max - ch_min)这种方法把每个通道自己的最小值映射到黑最大值映射到白适合看纹理和激活区域的位置。但有一个误区逐通道归一化会把“响应强度差异”抹平。如果两个通道一个有强烈激活、一个接近全零归一化后看起来都差不多亮。你想对比通道间的响应强弱应该用全局 min-max 或者保存原始 min、max、mean、std 数值画图只作为辅助。我通常会先打印这几个统计量再画图print(shape:, feat.shape) print(min:, feat.min(), max:, feat.max(), mean:, feat.mean())不要只看图。特征图可视化很容易骗人一个全零通道经过 min-max 归一化后可能显示成一片均匀灰色看着有点像有效特征实际上没有信息量。看了数值之后你才能判断这个通道是真的有响应还是根本“死了”。4.3 多通道特征图拼接和保存单个通道只是很小的矩阵直接画出来看不清。常用做法是把多个通道按网格摆成一张大图。可以用 matplotlib 画子图也可以用 torchvision.utils.make_grid 做网格图。先写一个可复用的保存函数import math import matplotlib.pyplot as plt import torch.nn.functional as F def save_feature_map(tensor, save_path, cols8, max_channels64, upscale_sizeNone): tensor tensor[0].detach().cpu() # [C, H, W] if upscale_size is not None: tensor F.interpolate( tensor.unsqueeze(0), sizeupscale_size, modebilinear, align_cornersFalse ).squeeze(0) feat tensor.numpy() if max_channels is not None and feat.shape[0] max_channels: feat feat[:max_channels] c, h, w feat.shape rows math.ceil(c / cols) fig, axes plt.subplots(rows, cols, figsize(cols * 2, rows * 2)) axes axes.flatten() for i in range(cols * rows): if i c: img normalize_channel(feat[i]) axes[i].imshow(img, cmapgray) axes[i].axis(off) plt.tight_layout() plt.savefig(save_path, dpi150) plt.close() save_feature_map(feature_map[layer1], layer1_feat.png, cols8, upscale_size(112, 112))这个函数做了三件重要的事限制max_channels避免 512 个通道全画出来。对低分辨率特征图做临时上采样方便肉眼看比如 layer4 的 7x7 可以放大到 56x56。用cmapgray明确画成灰度图避免 matplotlib 默认 colormap 导致颜色误导。4.4 特征图的数值范围怎么判断不同层、不同模型、不同输入feature map 数值范围天然不一样。没有统一标准可以套用但根据我的经验有几个判断方向第一层卷积的输出通常应该能看出边缘、纹理和背景的差异。如果全是一个平面先检查输入图片是否被正确读取Normalize 的 mean/std 是否写错。中间层特征图会出现不少低响应通道这是正常现象。如果大量通道同时全零或全恒定说明模型可能没训练好或者输入和模型预期分布差太多。深层特征图分辨率低语义更抽象直接肉眼不总能看出明确形状不必强求“像原图”。这时候要看不同类别之间是否有差异而不是看单张图本身像什么。如果 loss 已经变成 NaN 或者模型参数异常特征图数值会经常出现极大值或 nan。这时先用torch.isnan(tensor).any()检查张量再考虑可视化。5. 进阶卷积核可视化与 TensorBoard 记录5.1 可视化卷积核权重特征图是“数据经过卷积后的响应”卷积核是“网络学到的模板”。前者是中间层输出后者是模型参数。调试中核可视化也很有用。以 ResNet18 的输入卷积为例权重 shape 是 [64, 3, 7, 7]64 个卷积核每个卷积核有 3 个通道对应 RGB 输入卷积核尺寸是 7x7。因为输入有三个颜色通道可以把它们组合成一张彩色小图weight model.conv1.weight.detach().cpu().numpy() # [64, 3, 7, 7] fig, axes plt.subplots(8, 8, figsize(8, 8)) for i in range(64): ax axes[i // 8][i % 8] kernel weight[i].transpose(1, 2, 0) # [7, 7, 3] vmin kernel.min() vmax kernel.max() kernel (kernel - vmin) / (vmax - vmin 1e-8) ax.imshow(kernel) ax.axis(off) plt.tight_layout() plt.savefig(conv1_kernels.png, dpi150) plt.close()注意并不是所有卷积核都能这样合成彩色。第一个卷积层输入是 RGB 时才能这么做内部卷积核的输入通道数是上一层的输出通道数通常不是 3需要把每个输入通道单独显示或者取一个通道做灰度显示。卷积核可视化和 feature map 可视化在代码上是两个路子一个是 model 的 weight一个是 forward 的 output不要混在一起。5.2 用 TensorBoard 记录中间层输出如果只是看单张图matplotlib 就够了。但你在训练过程中想看特征图随 step 的变化或者想在同一页面对比多个 epoch 的变化TensorBoard 会更合适。PyTorch 自带torch.utils.tensorboard需要先安装 tensorboardpip install tensorboard然后在线程训练脚本里写入from torch.utils.tensorboard import SummaryWriter from torchvision.utils import make_grid writer SummaryWriter(runs/vis_demo) def write_feature_map(name, tensor, global_step): tensor tensor[0].detach().cpu() grid make_grid(tensor.unsqueeze(1), normalizeTrue, nrow8) writer.add_image(name, grid, global_step) write_feature_map(layer1, feature_map[layer1], global_step0) write_feature_map(layer4, feature_map[layer4], global_step0) writer.close()这里用tensor.unsqueeze(1)把 [C, H, W] 变成 [C, 1, H, W]make_grid 才会把每个通道当成一张单通道图处理。normalizeTrue会对整张网格做归一化适合快速预览。最终在终端运行tensorboard --logdir runs/vis_demo浏览器打开输出地址就能看到。TensorBoard 的好处不只是图片。你还可以把每层输出的 min、max、mean、std 用add_scalar和add_histogram记录这样模型跑几个 epoch 之后可以看到特征值的分布是怎么演变的。这个信息比单次截图更有说服力。5.3 在训练循环里可视化要注意什么训练过程中把 feature map 写入 TensorBoard最大的问题是频率。如果每个 batch 都写 2 层特征图TensorBoard 文件会越来越大训练也会被 IO 拖慢。我的做法是固定一个很小的采样间隔比如每个 epoch 只记录前几个 step或者只在验证集上记录一次。另外在训练循环里记录时如果是带有 Dropout 和 BatchNorm 的模型建议在记录特征图时临时切换到 eval 模式记录完再切回 train否则特征图每次 forward 都可能不一样。还有一点记录前最好确认当前 batch 的标签是什么。如果只保存特征图但不记录输入图片和标签后期回看时很难还原这是在处理什么样本。可以把输入 image 和标签也写入 TensorBoard使用writer.add_image和writer.add_text作为特征图的上下文。6. 常见问题排查没输出、全黑、shape 不对怎么处理6.1 先按顺序排查特征图可视化这种需求单独看某个报错经常会误导人。比如你看到RuntimeError: cant convert CUDA tensor to numpy第一反应可能是环境问题其实根因是忘了.cpu()。所以我一般会按下面的顺序排查先确认 hook 有没有触发。在 hook 里加一行 print或者看feature_map字典是否有对应 key。再确认模型是不是真的走了一次 forward。如果模型只是被实例化没有被调用hook 不会执行。确认输入张量 shape 和 dtype。ResNet18 要求 [N, 3, H, W]dtype 是 float32。确认模型处于 eval 还是 train。train 模式下 BN 和 Dropout 会带来额外变化。确认可视化张量做没做 detach、cpu、numpy。最后看数值。先打印 min、max、mean、std再画图避免被归一化骗了。6.2 高频问题对照表下面是我在调试中经常遇到的问题和对应处理思路。现象常见原因处理思路feature_map 里没有 key模型没执行 forward确认是否调用了 model(input)hook 执行了但 output 是 tuple某些自定义模块返回多个值取 output[0] 或打印 output 结构保存图片全黑特征值全负数且未归一化做 min-max 归一化保存图片全白或纯色所有通道数值几乎恒定打印 min/max检查该层是否失效shape 不是预期大小输入尺寸错误检查 Resize、padding、strideGPU 环境下报 numpy 错误tensor 还在 GPU 上先.cpu()再.numpy()hook 被重复调用多次同一模块被注册多次保存 handle用完后 remove编辑器里能跑脚本里没输出模型分支提前 return在目标模块加 print 确认执行路径训练中特征图每次都不一样模型是 train 模式加model.eval()或记录训练/验证模式feature map 数值有 nan模型参数异常或输入有 nan检查输入数据、loss、学习率、权重初始化这张表不是全部情况但覆盖了绝大多数初看“奇怪”的问题。遇到表中没列的现象先打日志把模块名、输入 shape、输出 shape、输出数值范围都打印出来通常能找到规律。6.3 一些值得长期保留的习惯可视化中间层输出这件事做完一次实验不难难的是把它变成稳定的调试习惯。我自己留了三个习惯你可以参考所有可视化工具都封装成函数。save_feature_map、write_feature_map这种函数从第一天就独立维护后面换模型不会散落一地重复代码。每次看特征图先打印统计量。只画图不打印 min/max/mean/std很容易把无效特征当有效特征。单样本先跑通再多层、多 batch。不要一上来就在整个数据集上开可视化先拿一张图确认 hook、保存、查看闭环是完整的。如果你也遇到模型训练正常但输出结果奇怪的情况先不要急着改结构或调学习率。把第一个卷积层和最后一个卷积层的输出拉出来看一眼很多时候答案就在那几张特征图里。