TensorFlow中堆叠打乱张量后恢复原张量的实现方法

核心思路 #

要恢复原张量，关键是记录打乱操作的索引映射关系——先保存原始数据的索引，打乱索引后用它来重排数据；后续恢复时，通过逆索引映射把处理后的数据还原回原始顺序，再拆分即可。

具体实现步骤 #

1. 定义原始占位符a和b #

import tensorflow as tf

# 定义对应形状的占位符
a = tf.placeholder(tf.float32, shape=(32, 2048), name="Placeholder:0")
b = tf.placeholder(tf.float32, shape=(64, 2048), name="Placeholder:1")

2. 拼接张量并生成原始索引 #

这里以样本维度拼接为例（将a和b合并为(96, 2048)的张量），若需在新维度堆叠（得到(2, 96, 2048)），逻辑可类比调整：

# 拼接a和b得到总张量
combined = tf.concat([a, b], axis=0)  # shape=(96, 2048)

# 生成原始索引，对应combined中每个样本的初始位置
original_indices = tf.range(tf.shape(combined)[0])  # shape=(96,)

3. 打乱索引并重排张量 #

# 打乱原始索引，得到乱序后的索引序列
shuffled_indices = tf.random.shuffle(original_indices)

# 根据乱序索引重排总张量，得到打乱后的输入
shuffled_combined = tf.gather(combined, shuffled_indices)

4. 送入网络处理 #

假设network_process为你的自定义网络处理函数：

# 对打乱后的张量执行网络处理
processed_shuffled = network_process(shuffled_combined)

5. 恢复原始顺序并拆分回a和b #

利用逆排列索引还原原始顺序，再拆分回对应张量：

# 计算逆索引：inverse_indices[i] 表示原始索引i在乱序索引中的位置
inverse_indices = tf.math.invert_permutation(shuffled_indices)

# 用逆索引重排处理后的张量，恢复原始顺序
processed_original = tf.gather(processed_shuffled, inverse_indices)

# 拆分得到对应原张量的处理结果
processed_a = processed_original[:32, :]
processed_b = processed_original[32:, :]

关键说明 #

• 若使用tf.stack进行堆叠操作，只需调整维度顺序后再执行索引打乱逻辑，核心的索引记录与逆映射恢复方法通用。
• tf.math.invert_permutation是实现逆索引的核心函数，可直接生成排列的逆序映射，无需手动计算。

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