使用tf.stack堆叠数组后形状不符，如何调整为原RGB张量同形状？

嘿，我一眼就看出问题出在哪啦！你用tf.split拆分rgb张量后，得到的r、g、b每个都是形状为(?, 672, 672, 1)的4维张量——最后多了一个长度为1的通道维度。这时候你用tf.stack在第3轴堆叠，它会把这三个4D张量当成元素去堆叠，结果就额外多了一维，变成(?, 672, 672, 3, 1)，自然和原rgb的形状对不上。

给你两种简单的解决办法：

方法1：用tf.concat替代tf.stack（推荐） #

tf.concat是沿着已存在的维度拼接张量，不会新增维度，刚好符合你的需求：

import tensorflow as tf

rgb = tf.placeholder(tf.float32, [None, 672, 672, 3])
r, g, b = tf.split(rgb, 3, 3)
# 用concat在第3轴拼接三个单通道张量
bgr = tf.concat([b-23.5, g-30.6, r-40.1], axis=3)

print("rgb :", rgb)
print("r :", r)
print("bgr :", bgr)

运行后bgr的形状就会变成(?, 672, 672, 3)，和原rgb完全一致。

方法2：先去掉多余维度再用tf.stack #

如果一定要用tf.stack，可以先通过tf.squeeze把每个拆分后的张量的最后一维（长度为1的维度）去掉，再进行堆叠：

import tensorflow as tf

rgb = tf.placeholder(tf.float32, [None, 672, 672, 3])
r, g, b = tf.split(rgb, 3, 3)
# 移除每个张量的最后一维
r_squeezed = tf.squeeze(r, axis=3)
g_squeezed = tf.squeeze(g, axis=3)
b_squeezed = tf.squeeze(b, axis=3)
# 堆叠后会自动新增第3轴，得到目标形状
bgr = tf.stack([b_squeezed-23.5, g_squeezed-30.6, r_squeezed-40.1], axis=3)

print("rgb :", rgb)
print("r :", r)
print("bgr :", bgr)

额外说明 #

• tf.stack的作用是将多个同形状的低维张量组合成一个高维张量（比如3个形状为(A,B,C)的张量，stack后变成(3,A,B,C)或其他轴的组合）；
• tf.concat则是在已有维度上扩展，把多个张量在指定轴上“粘”在一起，不会改变张量的维度数。

这里显然第一种方法更直接高效，推荐使用~

内容的提问来源于stack exchange，提问作者spicy burrito