一开始看到draw_tensor2psd.py0126v1这个文件名的人多半会先愣一下这到底是个画张量的脚本还是个导出 PSD 的脚本。我用这个命名其实是为了把一个折腾了几周的小工具版本固定下来。简单说这个脚本解决的事情非常具体把深度学习模型中间层的张量按通道拆开导出成一份带分层的 PSD 文件。为什么不用 PNG 拼图为什么非要用 PSD以及从张量到图层的每一步里有哪些坑这篇文章把话说透。如果你也在做模型可视化、特征图调试或者需要把算法输出的掩码、热力图、深度图交给旁边工位的设计师继续加工那这篇笔记应该能帮你少走不少冤枉路。1. 为什么我放弃了 matplotlib 和 PNG非要导出 PSD1.1 特征图可视化的两种典型需求做深度学习的人对“特征图可视化”都不会陌生但你可别以为它只有一种用法。我自己的体会是可视化需求至少能分成两个方向。第一个方向是调试型。模型某个层的输出不对我想快速看每个通道长什么样确认是不是全变成零了或者某一通道明显捕获了边缘、纹理、语义信息。这种场景下用常见的画图库把多个通道拼成一张大网格图几行代码就完事效率高也没什么学习成本。第二个方向是协作交付型。你要把模型某个阶段的输出交给不写代码的设计师或产品同学去评估。对方拿到一张 PNG 拼接图顶多放大缩小、看个热闹没办法单独把某一个通道拿出来修一修、调一调透明度也没办法在同一坐标下对比多个通道的变化。这时候你需要的不是一张“结果图”而是一份“可编辑的分层文件”。我最初在用方图凑合了一段时间直到有次要给一组注意力特征图做对比每张图都涉及 16 个通道、5 个不同时间步。PNG 网格图拼出来以后想单独看第 7 个通道在时间步 3 和步 4 的差异只能反复切换图片眼睛都快看花了。也就是从那一刻起我开始认真考虑把张量直接落成分层 PSD。1.2 PSD 的图层、混合模式与文件解析优势PSD 文件格式最打动我的一个点是它本质上就是一个“容器”里面可以有图层、图层组、混合模式、透明度、蒙版。对于张量可视化来说这就是天然友好的结构。你可以把张量的通道维理解成“一组同一视角下的不同图像”。一个形状为[H, W, C]的张量其中 C 个通道完全可以一一对应 PSD 里的 C 个图层。打开文件后想关哪个图层就关哪个想调整某个通道的透明度就调甚至把几个通道用“变亮”“正片叠底”之类的混合模式叠在一起观察它们的交互关系。这比一次性把所有通道画在同一张图上灵活得多。还有一点容易被忽略PSD 并不一定要靠某款商业图像软件才能打开现在很多开源查看器、在线预览工具都兼容这个格式。你把分层文件发给对方对方不用安装特定软件也能正常查看和基础编辑。从兼容性角度来说PSD 作为一个“交付中间格式”比一张拼好的 PNG 要稳得多。很多人担心 PSD 文件写起来麻烦其实不需要自己手动解析二进制。我用的思路是用 Python 侧找现成的 PSD 写入能力把每个通道图像作为独立图层传进去最后统一保存。只是要注意不同能力的库对灰度图层、16 位深度的支持程度不一样这里我先卖个关子后面专门说踩坑。1.3 版本号 0126v1 的实际含义可能有人会问文件名叫draw_tensor2psd.py0126v1后面这串数字和字母是什么规律这里其实是个朴素的版本标记习惯。0126表示我最后一次大规模改动的日期是 1 月 26 日v1表示这是该天整理出的第一个稳定版本。为什么不叫final、final2、真的不改了因为这类名字过两周再看就是灾难。命名带日期和版本号至少能保证两点第一你知道当前这份代码是什么时候的状态第二当你把同一个脚本发给别人时对方可以根据版本号锁定行为差异。更关键的是可视化脚本这种工具看起来不起眼但输出结果一变后面所有人的判断都会跟着变。我用0126v1固定下来以后如果某个图层顺序或命名规则做了调整只会在新版本里体现旧版本依然可以复现对应的 PSD 结果。这也是我建议每个做算法工程的人养成的习惯玩具脚本也要有版本意识。2. draw_tensor2psd.py0126v1 的输入张量约定与图层映射方案2.1 输入格式统一CHW 和 BCHW 的入口处理写这个工具之前我第一件做的事情不是写代码而是先把“输入张量到底长什么样”这件事定死。模型中间层张量的形状通常有几种最常见的是[B, C, H, W]的四维张量不少可视化脚本希望我把第一维 Batch 里的某一张图挑出来也有直接拿到[C, H, W]的三维张量还有少部分情况下拿到的是[H, W, C]也就是已经排成图像格式的 Tensor。为了不让调用方在使用的时候还要去记“这里要转置”“那里要 squeeze”我在脚本入口处做了一个统一处理如果是四维张量[B, C, H, W]默认取 Batch 里的第一张或者通过参数指定某一张如果是三维[C, H, W]内部自动转成[H, W, C]如果是三维[H, W, C]直接进入下一步如果二维[H, W]把它当作单通道张量处理。统一成[H, W, C]之后后面所有的通道切片、图层命名、逐一写入 PSD 的逻辑就都完全一致了。这个“入口统一”看起来很简单实际能帮你省掉大量重复判断。2.2 通道到图层的三种映射模式拿到[H, W, C]之后怎么把 C 个通道变成 PSD 图层一开始我想得很简单一个通道一个图层但实际使用中很快就发现并不是所有模型输出都适合这种“全展开”的方式。我最后在脚本里做了三种映射模式按参数切换。第一种是“全部通道铺开”。适合通道数不多的情况比如 RGB 三通道、深度图单通道或者语义分割的 5 类概率图。每一层单独显示用户能非常清楚地看到每个通道的独立响应。第二种是“分组图层”。适合通道数很多的情况比如某个特征图有 64 个通道如果全部铺开PSD 的图层列表会变得非常恐怖。此时我会按通道索引分成若干组比如每 8 个通道合成一个图层组在 PSD 里也体现出清晰的层级。这个做法的好处是外部视觉上不会因为图层太多而卡顿同时想看细节时依然可以展开组。第三种是“预览图 细节层”。这种做法适合需要交付给非算法同学的情况。最上面放一张把所有通道加权求和、或者取最大值的预览图层下面再放各个原始通道图层。默认打开文件时只看到预览图想要深入了解细节再一层层打开看。既方便沟通也保留了灵活性。2.3 数值范围和归一化预处理模型输出张量的数值范围千奇百怪这应该是可视化脚本里最容易出问题的地方。有些模型的输出是[0, 1]区间比如经过 Sigmoid 的掩码有些是[-1, 1]区间比如扩散模型在去噪过程中经常出现还有一些是任意范围的 logits最小值和最大值可能相差非常大。如果直接把原始数值转成像素很多情况下看到的画面要么全黑、要么全白毫无信息量。所以要进入 PSD 前必须做归一化处理。我在这版脚本里采用的默认策略是“线性拉伸到 8 位灰度”公式大概是这样min_val tensor.min() max_val tensor.max() norm (tensor - min_val) / (max_val - min_val 1e-8) norm norm * 255.0但注意min和max如果分别取的是整张张量的全局极值细节表现力往往会受极端像素影响。所以我额外加了一个百分位裁剪的开关默认用[1, 99]百分位代替全局最小最大值这样能有效避免个别极端点把整体亮度拉偏。3. 核心实现四步把 Tensor 改成 PSD 分层文件3.1 第一步反归一化与数值裁剪如果你在调试阶段已经得到的是一个[0, 1]的浮点张量理论上可以直接乘 255 转成图像。但实际模型输出很少这么干净尤其在中间层经常出现负值或绝对值很大的数字。我在脚本的第一步里做了两件事先把输入从任意范围拉到可显示范围再限制到合法的像素区间。反归一化并不是要还原训练时的“原始输入图像”而是要把模型内部响应变成人眼可感知的亮度映射。如果你希望解读的是“相对激活强弱”而不是“绝对像素颜色”用线性拉伸是很合理的。这步需要注意浮点精度。float32在归一化之后往往是带小数的转成 8 位整数前要执行clip(0, 255)。如果直接把负数或大于 255 的数值传给图像写入层有的库会静默截断有的库会直接报错两种行为都会影响你定位问题。3.2 第二步维度重排、切片和通道分离因为入口已经统一成[H, W, C]切片逻辑就非常简单了。如果你选择“全部通道铺开”就是对最后一维循环每次取出tensor[:, :, i]得到一个单通道矩阵。如果选择“分组图层”就先计算每组包含哪几个通道索引再逐个取出来统一挂到同一个组名下面。在做切片的同时我会顺手记录这个通道的统计信息包括均值、标准差、最小值、最大值。这些信息后续会放进图层名或者图层注释里方便你在 PSD 中快速判断某个通道是否激活异常。比如一个通道的标准差接近 0基本可以断定它“死”了这时候不必打开图层光看名字就能发现。3.3 第三步逐图层缩放并写入 PSD通道切片完成之后需要把二维矩阵变成“PSD 能识别的图层数据”。最直接的做法是把每个二维矩阵包装成灰度图像对象然后递增地添加到 PSD 文档中。我采用的顺序是从输出列表的最后往前放图层因为 PSD 的视觉顺序通常是最下面的图层在最底部第一层会在最上面。如果你的通道顺序本身有语义含义比如第 0 个通道是背景、第 1 个通道是物体保持“从 0 到 C-1 从下往上放”的规则在 PSD 里也就自然形成对应关系。如果某些通道需要在视觉上突出我还会在写入时设置图层不透明度。比如默认所有通道都是 100% 不透明但预览图可以设成 60%-80%这样叠加查看时不会完全遮盖下面的通道。3.4 第四步图层命名、顺序和元数据图层命名这件事一开始我没有太在意直接让脚本用“channel_0”“channel_1”这样的名字。直到有一次导出了一个 128 通道的特征图打开 PSD 后图层列表全是一串序号我才意识到命名如果没有信息量这个文件基本没法用。现在我的命名规则包括三层信息层级索引、通道在原始张量中的编号、通道统计特征。比如L2_ch07_mean0.12_max0.98一眼就能看出这是第 2 组特征图的第 7 个通道整体响应偏弱但最大值不低可能是一个只在局部有强激活的通道。如果你是在做扩散模型的可解释性还可以把时间步 t 编进图层组名例如step50_block2导出后按步数展开不同时间步之间的对比会非常清晰。3.5 代码骨架一版可运行的最小实现这里给出一份精简后的核心逻辑用到了我自己封装的 PSD 写入模块函数名保留成psd_writer你用自己习惯的 PSD 库替换即可。# draw_tensor2psd.py0126v1 核心逻辑骨架 import numpy as np from psd_writer import PSDDocument, Layer def prepare_tensor(tensor, batch_index0): 统一输入为 [H, W, C] 的 float32 数组。 四维张量默认取 batch_index 这一张。 arr np.asarray(tensor, dtypenp.float32) if arr.ndim 4: arr arr[batch_index] # [C, H, W] if arr.ndim 3 and arr.shape[0] 16: # 这里只在形状明显是 C,H,W 时做转置 if arr.shape[0] arr.shape[1] and arr.shape[0] arr.shape[2]: arr np.transpose(arr, (1, 2, 0)) if arr.ndim 2: arr arr[:, :, np.newaxis] return arr def normalize(arr, lower_percentile1, upper_percentile99): lo np.percentile(arr, lower_percentile) hi np.percentile(arr, upper_percentile) clipped np.clip(arr, lo, hi) normalized (clipped - lo) / (hi - lo 1e-8) return (normalized * 255.0).astype(np.uint8) def tensor_to_psd(tensor, output_path, batch_index0, channel_group_sizeNone): arr prepare_tensor(tensor, batch_indexbatch_index) h, w, c arr.shape doc PSDDocument(widthw, heighth) if channel_group_size is None: for i in range(c): channel_img normalize(arr[:, :, i]) layer Layer(namefch{i:03d}, imagechannel_img) doc.add_layer(layer) else: for group_id in range(0, c, channel_group_size): group_name fgroup_{group_id:03d} doc.add_group(group_name) for i in range(group_id, min(group_id channel_group_size, c)): channel_img normalize(arr[:, :, i]) layer Layer(namefch{i:03d}, imagechannel_img) doc.add_layer_to_group(layer, group_name) doc.save(output_path) return output_path这份代码虽然不能直接在你的环境里跑通但核心流程就是上述四步统一形状、切通道、归一化、写入 PSD。真实项目里你要做的就是把psd_writer换成自己环境里可用的 PSD 写入库并处理好类型转换。4. 实测中遇到的几个“阴间问题”和对应解法4.1 灰度图层被识别成索引色颜色全部失真第一次导出 32 通道特征图时我打开 PSD 后发现一个诡异现象部分图层显示正常部分是灰蒙蒙的而且明明每个通道都是独立灰度图可几个图层叠加后完全不是预想的亮度叠加效果。排查到最后问题出在灰度图层的存储模式上。某些写入库在处理单通道图像时会偷懒写成索引色或位图模式而很多 PSD 查看器对这类图层模式的兼容性并不好导致渲染结果和预期完全不一致。解决方案也很简单不要让写入库自动猜测单通道怎么存而是在创建图层时显式指定为 RGB 模式并把灰度值复制到 RGB 三个通道上。这样虽然文件会稍微大一点但兼容性稳定得多再也不需要担心打开后颜色错乱。4.2 16 位深度与 8 位深度的兼容性接着 4.1 深入说一个点。有的模型输出范围特别宽比如扩散模型中间层特征图的数值可能从 -40 到 40为了保留更细腻的层级差异我很自然想到用 16 位深度保存 PSD。结果试了一圈发现16 位通道在不少查看器里支持并不好有的能打开但图层叠加效果不对有的干脆只能看到预览图。后来我学到一个经验可视化交付场景优先用 8 位深度因为本质上是给人看相对分布不是做像素级科学测量。如果你确实需要做定量对比就把归一化前的统计参数单独存一份 JSON而不是把希望寄托在 16 位 PSD 上。8 位 正确归一化已经能覆盖绝大多数需求。4.3 图层名里的中文和特殊符号模型层名经常带着点号、斜杠、冒号比如block.1.attn.proj。如果直接用这个名字去创建 PSD 图层部分写入库在保存时可能会截断或者替换非法字符导致图层名对不上号尤其跨系统传递后再打开问题更明显。我的做法是加一个清洗函数把所有非字母数字的字符统一替换成下划线同时允许我用自定义别名覆盖原始层名。比如在配置里写清楚block.1.attn.proj - attn_proj这样脚本在导出时既能保留人话命名又不会触发写入库的字符限制。4.4 大批量导出时的内存失控与临时文件通道数多的场景下最危险的还不是 PSD 写入本身而是你在写入前把几十张、上百张图像对象全部创建好放在内存里。我踩过一个大坑导出一次 132 个通道、分辨率 1024×1024 的特征图脚本内存直接涨到 7GB 多然后被杀掉。后来发现我对每个通道都先转成 RGB 三通道图像再放到一个 list 里等最后一起提交这个“等最后一起提交”就是内存爆炸的根源。正确姿势是流式写入每处理完一个通道立刻写入 PSD 文档并释放局部变量。如果 PS 写入库不支持流式至少要用生成器分批次处理或者把中间结果先写到磁盘上的临时 PNG 目录最后再统一合成 PSD。内存峰值能降一个数量级。4.5 浮点张量里出现 NaN 和 Inf 怎么办最后一个常见问题是张量里混进 NaN 或者 Inf。归一化的时候只要有一个 NaNmin和max就全变 NaN整张图导出结果就是黑的或者花的而且这种问题往往不是可视化脚本引起的是上游模型或数据问题。我在脚本里做了一个显式开关默认检测到 NaN/Inf 时把异常值替换成当前数组的有限值范围内的最小值或最大值并打印警告。这样既不会让整个导出直接失败也能通过图层名里的统计特征一眼看出这个通道有异常。对于调试工作流来说把“模型出问题”和“可视化脚本出问题”分开处理能节省大量定位时间。5. 这个脚本真正有用的三个使用场景5.1 扩散模型注意力特征图的可解释性分析如果你研究过扩散模型应该知道它的 UNet 中间层会输出很多特征张量这些张量在时间步推进过程中不断变化。用 PSD 分层导出之后你可以把不同时间步的输出按图层组存放然后逐个开启对比观察某个空间区域的特征是何时开始形成的。我试过把去噪过程中间步的特征图导出成分层 PSD再用混合模式把几个通道叠在一起看确实能看到一些有趣的现象比如某些通道负责结构轮廓某些通道负责材质纹理它们在去噪早期和晚期的激活模式有很明显差异。这些观察如果靠 PNG 拼图几乎很难形成直觉。5.2 语义分割和深度估计的掩码微调语义分割模型输出的概率图也特别适合转成 PSD。每个类别对应一个图层概率越高图层越亮配合透明度和颜色映射可以直接看出模型在不同类别边界上的犹豫程度。深度估计任务更直接输出的是一张单通道深度图。用这个脚本导成 PSD 后设计师可以把深度图作为一个图层叠加到底图上做景深效果。图层透明度一拖算法结果和艺术表达之间的缝隙就被填上了。5.3 跨岗位协作模型结果交给设计师二次处理这也是最让我觉得这个脚本有价值的地方。之前跨岗位协作时我把特征图拼成 PNG 发给设计师对方只能“看”没法“改”。现在导出分好层的 PSD设计师直接把某一层的对比度提亮一点或者把某个通道的混合模式改一下几分钟内就能验证一个表达效果并且改完还能把结果再发回给我大家讨论的锚点也从抽象的数字变成了具体的画面。如果你所在团队的算法和设计协作比较频繁我建议你也试试这个思路。哪怕只是简单地把模型输出的几个图分层导出双方的沟通成本都会降一截。6. 给后来者的一点点配置建议和迭代方向6.1 环境依赖与版本锁定这种小工具最容易被忽略的就是环境复现。你可能今天在笔记本上跑通了下个月换台机器安装的依赖版本不同PSD 导出结果就变了。我在项目里专门加了一个依赖清单把 Python 版本、核心数值库版本、PSD 写入库版本全部固定住。每次大版本迭代后都会重新导出同一份测试张量对比前后文件哈希是否一致。听起来有点小题大做但可视化脚本一旦成了团队协作的工具稳定复现就是第一优先级。6.2 面向大图的图层懒加载思路现在遇到的另一个瓶颈是超大分辨率张量。未来如果要导出 4096×4096 的多通道特征图即使内存能扛住写入 PSD 的时间也会非常久。我计划中的改进方向是引入“懒加载”思路不是把每个通道的全部像素一次性打包进 PSD而是先写入低分辨率预览图保留配置文件记录原始数据的裁剪方式、存储路径。当用户在 PSD 里确实需要查看原始分辨率细节时再通过外部引用或后续替换高分辨率图层完成。这个方案还在验证阶段但对于超大图可视化它应该比硬扛内存更可持续。6.3 后续可以扩展的功能点最后说一下这个脚本可以继续长出来的新功能。第一个是命令行参数化通过一个配置文件指定输入张量的路径、图层分组策略、通道命名别名、是否做百分位归一化这样不同项目之间不需要改代码只改配置就行。第二个是批量导出能力。很多实验会同时产生几十份中间张量如果脚本只支持单文件导出后续还得自己写循环不优雅。批量模式还能顺带生成一份索引表方便你从实验记录反向定位到具体 PSD 文件。第三个是“对比模式”。给定两份同形状的张量脚本自动生成“A 图层、B 图层、差异图层”三组结构差异图层直接标出每像素绝对差。这个功能对 A/B 测试特别实用我的计划是把它做到draw_tensor2psd的下一个版本里。我自己用得越久越觉得这种工具的价值不在于代码有多精巧而在于它真的能把模型排查和创意协作这两件看起来不太相关的事用一份分层文件串起来。如果你也准备复制这个思路我的建议是先只支持一种输入格式把所有通道按最简单的方式铺开跑通一遍再接分组、接归一化选项。等流程稳定了再考虑性能优化和更多使用场景别一上来就想一口吃成胖子。