看到 draw_tensor2psd.py 这个文件名熟悉深度学习模型落地流程的朋友应该能秒懂这是要把张量tensor直接画成PSD文件。我在某个周五下午被设计师一句这个能直接给我PSD吗问住之后才动手写了这个脚本。版本号0126v1说白了就是1月26日改的第一版后面还有0127v2、0203v3这些乱七八糟的迭代文件名越改越长工具倒是越来越顺手。这篇文章就围绕这个脚本展开说说我为什么要把tensor转成PSD里面踩过哪些坑以及一份能直接跑起来的完整实现。适合有三四个月Python基础、正在做模型输出可视化或者经常要跟设计师打交道的算法工程师参考。1. 为什么要做Tensor转PSD这件事1.1 特征图不是给人看的但甲方和设计师要模型推理完你手上拿到的是一堆float32的张量shape可能是 [batch, channel, height, width]也可能是一个二维的注意力矩阵。这些数据在numpy里是精确的在训练日志里是能画曲线的但你要是直接把数组丢给设计师对方肯定是一脸茫然。真实需求往往长这样语义分割模型跑完一张街景图输出了19个类别的概率图或者Stable Diffusion的某个中间层特征图想拿来debug——这些都需要可视化。可可视化如果只是存一张PNG设计师后续想在原图上叠加、调整透明度、单独修某一类别的区域就得重新来一遍。PSD的好处在于它是带图层结构的我输出的每一个channel、每一张特征图都是一层对方拿到手可以直接打开改。这个脚本解决的痛点就是把模型输出的张量和设计工具里的可编辑文件之间的最后一公里打通。1.2 PSD的分层能力是刚需不是玄学很多人会问PNG不也能带透明度吗为什么非要PSD关键在可编辑性。PNG是一张拍扁的图所有像素烧在一个平面上。PSD本质是一个容器里面可以塞任意多个图层每个图层可以有自己的混合模式normal、multiply、screen等、不透明度、蒙版信息。比如我想把4张不同形状的mask叠在一张原图上用PNG我得先合成好、然后导出用PSD可以直接输出4个图层设计师在Photoshop里手动开关图层、调透明度十秒钟就能出一张对比图。另一个场景是debug。模型中间层特征图往往有几十上百个channel我通常会把每个channel或每组channel变成独立图层再配上不透明度降低一些的叠加模式。这样一眼就能看出哪个channel学到了边缘、哪个channel学到了纹理比打印一堆灰度图效率高得多。2. 数据管线的设计从Tensor到像素再到图层2.1 先搞清楚Tensor的形状与数值范围写脚本的第一步不是写代码是先明确输入。常见的tensor来源有PyTorch的torch.Tensor拿回来先.detach().cpu().numpy()TensorFlow的EagerTensor直接.numpy()已经导出的.npy文件这是最稳的格式拿到numpy数组之后首先要处理维度问题。我归纳了一下实际项目里无非这几种情况2D[H, W]单张灰度图最常见3D[C, H, W]多通道特征图每个通道单独成层4D[N, C, H, W]带batch的一般取data[0]即可除非你想把整个batch输出成一组图层数值范围这一关特别容易被忽略。模型输出可能是任意浮点范围比如 logits 里最大值12.8、最小值-7.3如果你直接做astype(np.uint8)所有负值都会被截成0大于255的直接变255结果就是一张花屏。正确做法是先归一化我用的是百分位裁剪加min-max缩放这个策略对付长尾分布的tensor最稳。2.2 归一化与伪彩色像素化不是直接转类型归一化这一步值得单独展开。很多初学朋友以为转成图片就是np.uint8(array)实际上至少要经过三步去掉无效值检查np.isfinite把NaN和inf排除掉否则后面统计min/max时会得到一堆垃圾结果。百分位裁剪我一般取1%和99%分位数作为上下界有效防止个别极端像素值把整体亮度拉低。min-max归一到 [0,1]再乘255转uint8。这是一个完整的函数def normalize_for_display(x): x np.asarray(x, dtypenp.float32) mask np.isfinite(x) if not mask.any(): return np.zeros_like(x, dtypenp.uint8) lo, hi np.percentile(x[mask], [1, 99]) if hi - lo 1e-8: hi lo 1e-8 y (x - lo) / (hi - lo) y np.clip(y, 0.0, 1.0) return (y * 255.0).astype(np.uint8)归一化之后灰度图直接可以当alpha通道或亮度层用。但很多时候光有灰度不够比如注意力热力图大家习惯用蓝→青→绿→黄→红这种jet伪彩色来表现强弱。这时候我建议提前生成一张256色的LUT用查表替代运行时调matplotlib速度能快一个量级。# 提前生成或者预置的jet LUTshape 为 (256, 3) JET_LUT np.load(jet_lut.npy) def apply_colormap(gray_u8): return JET_LUT[gray_u8] # 结果 shape [H, W, 3]2.3 图层映射策略一个Channel一个图层还是合并RGB图层怎么映射取决于你要表达什么。如果是语义分割结果每个类别对应一个二维概率图。我倾向于把每个类别作为独立图层图层名直接用类别名比如 road、building、sky。设计师可以单独隐藏某个类别或者把某个类别的混合模式从normal改成multiply做出来的效果非常灵活。如果是中间层特征图channel数量往往几十上百个。这时候全部分成一个图层PSD文件会巨大Photoshop打开也吃力。我通常的做法是先做通道分组把4个channel合成一张RGBA图或者把相邻channel的均值合成一张特征图。也就是说默认情况下我不把所有channel都倒进PSD而是提供一个stride参数隔几个channel取一层先看全局再深入细节。哪种情况适合把channel拼成RGB比如最终要给人看的合成效果图三个channel直接作为R、G、B通道拼一张底图再叠加其他mask图层。这样文件体积小视觉上信息密度也高。2.4 为什么不用psd-tools而选了pytoshopPython生态里能处理PSD的库不算多主流就两个psd-tools和pytoshop。psd-tools的强项是读取解析PSD里的图层、蒙版、文本对象对文件做无损深挖非常厉害但写多图层PSD的能力比较弱更多是把PIL图像叠成一个单一背景图层再保存。我之前试过用它从零创建多个带独立混合模式的图层绕了半天还是回到了二进制层面。pytoshop是MATTtt等作者维护的纯Python库定位就是读写PSD底层结构。它的API长得很像底层格式写每一个图层时需要显式地添加R、G、B、Alpha四个通道的二维数组。上手刚开始会觉得啰嗦但自由度够我能精确控制图层顺序、混合模式和不透明度这才是把tensor转成可编辑PSD该有的样子。如果你问我为什么不手写PSD二进制我只能说PSD的头部信息、通道压缩方式、图层蒙版结构这些细节看着简单真自己拼起来一天时间就没了而且每个版本Photoshop都会踩兼容性坑。有pytoshop这种轮子没必要再造。3. 动手实现draw_tensor2psd.py的完整实操3.1 环境准备依赖很简单四个库搞定pip install numpy pillow pytoshop如果你要加载PyTorch的tensor再装torch如果要用matplotlib的cmap做伪彩色临时调试时可以装但最终脚本里建议用LUT查表省得引入一个重依赖。我建议建个干净虚拟环境再装避免把系统Python搞乱。py3.8到3.12都能跑pytoshop在3.10和3.11下我都实测过没问题。3.2 核心函数tensor_to_psd核心代码不复杂但每个细节都有讲究。下面是一份精简过、能直接跑通的骨架import numpy as np import pytoshop from pytoshop import user, enums def normalize_for_display(x): x np.asarray(x, dtypenp.float32) mask np.isfinite(x) if not mask.any(): return np.zeros_like(x, dtypenp.uint8) lo, hi np.percentile(x[mask], [1, 99]) if hi - lo 1e-8: hi lo 1e-8 y (x - lo) / (hi - lo) y np.clip(y, 0.0, 1.0) return (y * 255.0).astype(np.uint8) def tensor_to_psd(data, save_path, namesNone, stride1): data np.asarray(data) # 统一成 [C, H, W] if data.ndim 4: data data[0] if data.ndim 3: c, h, w data.shape elif data.ndim 2: c, h, w 1, data.shape[0], data.shape[1] else: raise ValueError(f不支持的shape: {data.shape}) psd user.PsdFile() psd.width w psd.height h psd.nchannels 3 psd.ndepth 8 psd.colormode enums.ColorMode.rgb layer_count 0 for i in range(0, c, stride): gray normalize_for_display(data[i]) rgb apply_colormap(gray) # [H, W, 3] uint8 alpha gray.copy() # 黑色区域半透明可选 layer user.LayerInfo() if names is not None and i len(names): layer.name names[i] else: layer.name fchannel_{i:04d} layer.blend_mode enums.BlendMode.normal layer.opacity 255 layer.top 0 layer.left 0 layer.bottom h layer.right w layer.layer_type enums.LayerType.pixel layer.add_channel(-1, alpha) # alpha通道 layer.add_channel(0, rgb[:, :, 0]) layer.add_channel(1, rgb[:, :, 1]) layer.add_channel(2, rgb[:, :, 2]) psd.layers.append(layer) layer_count 1 psd.write(save_path) print(f已写出 {layer_count} 个图层 - {save_path})需要说明的是psd.layers.append和psd.layers.insert(0, ...)会直接影响Photoshop里图层面板的上下顺序。我踩过一次坑同一个test.psd用append生成后打开Photoshop发现图层顺序跟我想的完全反了——最想放上面的层跑到了最底部。所以如果你发现顺序不对换insert(0, layer)再跑一次就行半天就能定位。3.3 完整脚本与命令行入口光有核心函数不够脚本要能直接被调用我加了一个很朴素的命令行入口if __name__ __main__: import argparse parser argparse.ArgumentParser(description把tensor/npy导出成多图层PSD) parser.add_argument(--input, requiredTrue, help输入.npy文件路径) parser.add_argument(--output, defaultout.psd, help输出PSD路径) parser.add_argument(--names, defaultNone, help逗号分隔的图层名) parser.add_argument(--stride, typeint, default1, help每隔几个channel输出一层) args parser.parse_args() raw np.load(args.input, allow_pickleTrue) names args.names.split(,) if args.names else None tensor_to_psd(raw, args.output, namesnames, strideargs.stride)实际使用是python draw_tensor2psd.py --input seg_logits.npy --output seg.psd --names road,building,sky --stride 2这版代码在pytoshop 0.2.x上测试过如果你装的是更高版本个别属性名可能会有变化比如PsdFile.nchannels在有些分支里叫n_channels。遇到报错别慌先python -c import pytoshop, inspect; print(inspect.signature(pytoshop.user.PsdFile))看一下当前版本签名按实际情况微调就好。4. 实测排坑那些文档里不会写的问题4.1 输出全黑或全白九成是归一化的问题这是我被问过最多的问题没有之一。新手拿float张量直接astype(np.uint8)结果就是值域在[0,1]的浮点数全部变成0图片全黑值域在[0,255]的浮点数但带负数负数被截断图片有黑斑值域很大的logits超过255的直接白色细节全丢解决路径很简单就用上面normalize_for_display里那套先np.isfinite过滤再做1%-99%百分位裁剪最后min-max。如果你发现某个tensor的分布特别集中比如99%的值都在0.1附近可以把百分位从[1, 99]改成[0.5, 99.5]或者干脆用min/max试试肉眼调到满意为止。4.2 Photoshop打开后图层缩略图是灰色或提示损坏这个坑十有八九出在通道数据没写全。pytoshop创建RGB模式的PSD时每一个像素图层至少要包含R、G、B三个通道如果你想带透明再加一个索引为-1的alpha通道。我见过有人只写了R通道就保存Photoshop能打开但缩略图完全错乱。另外要注意通道数据的dtype必须是uint8把float数组塞进去会报错或生成无效文件。我的习惯是先np.ascontiguousarray(rgb[:, :, 0], dtypenp.uint8)再传给add_channel确保内存连续。4.3 大尺寸Tensor的内存与耗时问题某次我跑一个[256, 512, 512]的特征图每个channel都是512x512直接用上面脚本生成256个图层最后PSD文件接近2GBPhotoshop打开花了五分钟。原因是没做任何降采样和通道抽样。解决思路有两个一是用stride参数抽样每隔4个或8个channel输出一层先看整体二是对单个二维矩阵做下采样比如用PIL的resize((w//2, h//2), Image.BILINEAR)肉眼对比特征模式完全够用。内存方面尽量避免把整批数据一次性astype成多个副本。我在脚本里每处理完一个channel就马上把它组装进LayerInfo然后用del gray, rgb, alpha释放引用让Python的GC及时回收。4.4 pytoshop版本不同导致API不兼容pytoshop这个库迭代不算频繁但每次大版本都可能改属性名。比如我最早用的是0.1.xPsdFile里面直接叫width、height后来0.2.x加了nchannels、ndepth这些字段。如果你的脚本报AttributeError: PsdFile object has no attribute nchannels不一定是你写错了很可能是库版本差异。我的建议是不管代码从哪抄的先写一个6x6像素的单图层PSD再用pytoshop把它读回来验证确认能读能写再处理真实数据。这个最小闭环能节省至少半小时的排查时间。4.5 给TensorFlow/PyTorch用户的一个提醒采样与随机种子最后提一个跟画图无关但很容易踩的坑。如果你在把tensor导出之前用过torch.rand、torch.multinomial这类采样操作——比如扩散模型里根据logits采样类别——那么每次运行脚本导出的图可能都不一样因为随机种子没固定。要复现就在调用脚本前固定import torch torch.manual_seed(0)TensorFlow那边则是tf.random.set_seed(0)。千万别以为同一个模型同一个输入就会得到同一个输出GPU上的非确定性操作比你想象的常见。5. 一点个人体会写这个脚本最大的收获不是学会了pytoshop的API而是意识到可视化这件事工具选型很重要但更关键的是搞清楚数据最终会被谁、以什么方式使用。PSD作为一种中间产物它的价值恰恰在于不把信息焊死在位图里而是保留图层这种可操作的结构。我后来在这个基础上加了几个小功能比如输出带混合模式的叠加层、给每个图层配一个调好的不透明度、批量把一个batch的所有tensor导出成多个PSD。一开始的draw_tensor2psd.py0126v1早就进化成一套内部工具了不过核心思路没变让tensor以人可以二次加工的形式流出模型省掉一遍遍导TXT、导CSV、导PNG的重复劳动。如果你也要做类似的事我的建议很简单先用最小尺寸的tensor跑通链路确认PSD能在Photoshop里正常打开、图层顺序符合预期再往真实数据上堆。这个脚本不长但它能帮你省下的是每次对接设计时那半小时的导出-截图-重新绘制时间。