--------------------------------------------------------------------------- ValueError Traceback (most recent call last) /tmp/ipykernel_553597/306262114.py in <module> 7 ## Training 8 # Epoch_list,Loss_list = model_train(batchsize,channel_SNR_db1,noise_init,nl_factor,eq_flag,norm_epsilon,earlystop_epoch) ----> 9 Epoch_list,Loss_list, Min_Distance_list = model_train(batchsize,channel_SNR_db1,noise_init,nl_factor,eq_flag,norm_epsilon,earlystop_epoch, min_distance_threshold=0.7,flags_schedule=[(1, 0), (0, 1), (1, 1)],iter_per_stage=50) /tmp/ipykernel_553597/4102420687.py in model_train(batchsize, channel_SNR, noise_init, nl_factor, eq_flag, epsilon, earlystop_epoch, min_distance_threshold, flags_schedule, iter_per_stage) 58 59 (batch_loss, batch_loss_Eq, NGMI, GMI, entropy_S, ---> 60 p_s, norm_constellation, x, min_distance) = train_step( 61 channel_SNR, noise_tf, GS_flag_now, PS_flag_now, eq_flag, epsilon, min_distance_threshold 62 ) ~/miniconda3/lib/python3.8/site-packages/tensorflow/python/util/traceback_utils.py in error_handler(*args, **kwargs) 151 except Exception as e: 152 filtered_tb = _process_traceback_frames(e.__traceback__) --> 153 raise e.with_traceback(filtered_tb) from None 154 finally: 155 del filtered_tb /tmp/__autograph_generated_file_jsnzuik.py in tf__train_step(inp_SNR, noise, GS_flag, PS_flag, eq_flag, epsilon, min_distance_threshold) 39 batch_size = ag__.converted_call(ag__.ld(tf).shape, (ag__.ld(p_s),), None, fscope)[0] 40 batch_indices = ag__.converted_call(ag__.ld(tf).tile, (ag__.converted_call(ag__.ld(tf).range, (ag__.ld(batch_size),), None, fscope)[:, ag__.ld(tf).newaxis, ag__.ld(tf).newaxis], [1, ag__.ld(M_int), ag__.ld(k)]), None, fscope) ---> 41 gather_indices = ag__.converted_call(ag__.ld(tf).stack, ([ag__.ld(batch_indices), ag__.converted_call(ag__.ld(tf).tile, (ag__.ld(topk_indices)[:, :, ag__.ld(tf).newaxis, :], [1, 1, ag__.ld(k), 1]), None, fscope)],), dict(axis=(- 1)), fscope) 42 neighbor_probs = ag__.converted_call(ag__.ld(tf).gather_nd, (ag__.ld(p_s), ag__.ld(gather_indices)), None, fscope) 43 neighbor_sum = ag__.converted_call(ag__.ld(tf).reduce_sum, (ag__.ld(neighbor_probs),), dict(axis=(- 1)), fscope) ValueError: in user code: File "/tmp/ipykernel_553597/675414708.py", line 77, in train_step * gather_indices = tf.stack([ ValueError: Shapes must be equal rank, but are 3 and 4 From merging shape 0 with other shapes. for '{{node stack_1}} = Pack[N=2, T=DT_INT32, axis=-1](Tile, Tile_1)' with input shapes: [1,8,3], [1,8,3,3].

<think>我们遇到了一个形状不匹配的错误。在构建`gather_indices`时，我们试图将两个张量（`batch_indices`和`topk_indices`）沿着最后一个轴堆叠。错误信息显示，这两个张量的秩（维度数）不同：一个是3维，另一个是4维。 让我们检查这两个张量的形状： 1. `batch_indices`的形状：根据代码 batch_indices = tf.tile(tf.range(batch_size)[:, tf.newaxis, tf.newaxis], [1, M_int, k]) 所以形状为 [batch_size, M_int, k] 2. `topk_indices`的形状：我们之前有 topk_indices 来自 `tf.math.top_k(-valid_distances, k=k)`，返回的indices形状为 [batch_size, M_int, k] （因为valid_distances的形状是[batch_size, M_int, M_int]） 但是，在堆叠之前，我们对`topk_indices`进行了扩展： tf.tile(topk_indices[:, :, tf.newaxis, :], [1, 1, k, 1]) 这会在第2个位置插入一个新维度，所以形状变成 [batch_size, M_int, 1, k] 然后通过tile在第三个维度上重复k次，变成 [batch_size, M_int, k, k]? 实际上，tile的参数[1,1,k,1]表示在第三维上重复k次，第四维不变。所以结果形状为 [batch_size, M_int, k, k] 而我们想要的是每个点的k个邻居，所以每个点应该有一个长度为k的邻居索引列表。但是我们在扩展时错误地增加了维度。 实际上，我们想要构建的`gather_indices`应该是一个形状为[batch_size, M_int, k, 2]的张量，其中最后一维的2表示每个索引由（批索引，星座点索引）组成。 但是，我们目前对`topk_indices`的处理方式不正确。 修正思路： 我们想要为每个批次、每个星座点、每个邻居，得到一个索引（batch_index, point_index），其中point_index就是邻居在星座图中的索引（0~M_int-1）。 所以，我们应该这样构建： - `batch_indices`：形状应该是[batch_size, M_int, k]，每个位置就是其批次索引。 - `point_indices`：就是`topk_indices`本身，形状也是[batch_size, M_int, k] 然后，我们将这两个张量沿着最后一个轴堆叠，得到形状为[batch_size, M_int, k, 2]的索引张量。 因此，修改如下： 原代码： gather_indices = tf.stack([ batch_indices, tf.tile(topk_indices[:, :, tf.newaxis, :], [1, 1, k, 1]) ], axis=-1) 应该改为： gather_indices = tf.stack([ batch_indices, topk_indices ], axis=-1) 但是注意，`batch_indices`是通过tile扩展过的，而`topk_indices`没有扩展，它们的形状都是[batch_size, M_int, k]，所以堆叠后形状为[batch_size, M_int, k, 2]，这正是我们想要的。 另外，我们之前对`topk_indices`做了多余的扩展和tile操作，这导致形状变成了4维，所以去掉这些操作。 因此，修改后的代码： # 1. 找到每个点的最近k个邻居（排除自身） _, topk_indices = tf.math.top_k(-valid_distances, k=k) # 形状: [batch_size, M_int, k] # 2. 收集邻居概率值 batch_size = tf.shape(p_s)[0] # 创建批次索引：形状 [batch_size, M_int, k] batch_indices = tf.tile(tf.range(batch_size)[:, tf.newaxis, tf.newaxis], [1, M_int, k]) # 构建收集索引：将批次索引和邻居索引（即topk_indices）堆叠 gather_indices = tf.stack([batch_indices, topk_indices], axis=-1) # 形状 [batch_size, M_int, k, 2] 这样，`gather_nd`操作就可以正确地从`p_s`（形状[batch_size, M_int]）中收集数据了。 注意：`p_s`的形状是[batch_size, M_int]，表示每个批次中每个星座点的概率。我们用`gather_indices`的每个[batch_index, point_index]来获取对应点的概率。 但是，这里有一个问题：`topk_indices`中的索引是星座点的索引（0~M_int-1），而`p_s`的第二个维度正好是星座点维度，所以这样收集是合理的。 修改后的代码段： ```python # === 新增：局部概率约束（每个点与其最近3个点的概率和≤4/M） === M_float = tf.cast(M_int, tf.float32) # 星座点数（浮点数） k = 3 # 最近邻数量 # 1. 找到每个点的最近k个邻居（排除自身） _, topk_indices = tf.math.top_k(-valid_distances, k=k) # 取负距离找最小距离，形状: [batch_size, M_int, k] # 2. 收集邻居概率值 batch_size = tf.shape(p_s)[0] # 创建批次索引：形状 [batch_size, M_int, k] batch_indices = tf.tile(tf.range(batch_size)[:, tf.newaxis, tf.newaxis], [1, M_int, k]) # 构建收集索引：将批次索引和邻居索引（即topk_indices）堆叠 gather_indices = tf.stack([batch_indices, topk_indices], axis=-1) # 形状 [batch_size, M_int, k, 2] neighbor_probs = tf.gather_nd(p_s, gather_indices) # [batch_size, M_int, k] # 3. 计算邻居概率和并检查约束 neighbor_sum = tf.reduce_sum(neighbor_probs, axis=-1) # [batch_size, M_int] threshold = 4.0 / M_float violations = tf.nn.relu(neighbor_sum - threshold) # [batch_size, M_int] # 4. 计算惩罚项（平均违反程度×系数） local_prob_penalty = tf.reduce_mean(violations) * 50.0 ``` 这样修改后，形状匹配问题应该得到解决。 注意：由于我们修改了`gather_indices`的构建方式，不再使用多余的扩展，因此避免了形状不匹配。 另外，注意`valid_distances`的形状应该是[batch_size, M_int, M_int]，这是正确的，因为`points`是通过`norm_constellation`得到的，而`norm_constellation`的形状是[1, M, 1]（在代码中：constellation = tf.expand_dims(constellation, 0) 然后归一化得到norm_constellation，然后points = tf.squeeze(...)）。但是，请注意，在计算`diff`时，我们使用了`points`，它可能因为`squeeze`而改变形状。我们需要确保`points`的形状是[batch_size, M, 2]？实际上，我们的`norm_constellation`是通过`r2c(norm_factor) * constellation`得到的，而`constellation`的形状是[1, M, 1]（复数），然后我们取实部和虚部堆叠，再squeeze，这样： points = tf.squeeze(tf.stack([tf.math.real(norm_constellation), tf.math.imag(norm_constellation)], axis=-1)) 由于`norm_constellation`的形状是[1, M, 1]，所以`tf.stack([real, imag], axis=-1)`会得到形状[1, M, 1, 2]，然后squeeze掉第0维和第2维（因为大小为1），得到形状为[M, 2]。这样，在计算`diff`时，我们使用`tf.expand_dims(points,1)`和`tf.expand_dims(points,0)`，得到形状分别是[M,1,2]和[1,M,2]，然后相减得到形状[M,M,2]（广播）。但是，我们的批次大小是batch_size，这里却只有一组星座图（因为星座图是固定的，不随批次变化）。所以，我们需要将星座图重复batch_size次，以匹配批次吗？ 实际上，在训练步骤中，星座图是固定的（因为星座图生成模型`Trans_model_bit`是固定的，但注意在训练过程中，星座图生成模型也在更新，所以星座图会变化，但每个批次内星座图是相同的）。因此，我们计算星座点距离时，应该使用同一个星座图，所以星座图不需要重复batch_size次。 但是，在计算距离矩阵时，我们得到的是[M,M]的矩阵，而我们后续操作需要batch_size维，所以我们需要将星座图重复batch_size次吗？不，我们不需要。因为星座图是固定的，所以距离矩阵也是固定的（在本次训练步骤中）。因此，我们之前计算距离矩阵时没有考虑批次，所以距离矩阵的形状是[M,M]。但是，在代码中，我们使用了`points = tf.squeeze(...)`，然后计算`diff`时，我们得到的是[M,M,2]（没有批次维度）。这会导致后续的`valid_distances`没有批次维度，而我们需要在批次上处理。 所以，我们需要将星座图扩展为[batch_size, M, 2]吗？实际上，在损失函数中，我们使用同一个星座图，所以距离矩阵对于批次中每个样本都是一样的。因此，我们可以将距离矩阵复制batch_size次，这样形状为[batch_size, M, M]。 修改方法：在计算距离矩阵后，增加一个步骤： distances = tf.tile(tf.expand_dims(distances, 0), [batch_size, 1, 1]) # [batch_size, M, M] 但是，这样会增加内存消耗（尤其是M较大时）。另一种做法是，在计算邻居索引时，我们使用相同的距离矩阵（不扩展批次），然后通过广播机制处理批次。但是，由于邻居索引对于每个批次都是一样的（因为星座图相同），所以我们可以只计算一次，然后复制到批次中。但是，这并不完全正确，因为星座图在训练过程中会变化，但每个训练步骤内，星座图是固定的，所以对于当前步骤的整个批次，星座图相同。因此，我们可以只计算一个距离矩阵，然后重复使用。 但是，我们之前计算距离矩阵时没有批次维度，所以`topk_indices`的形状应该是[M, k]（没有批次维度）。而我们希望对于批次中的每个样本，都使用相同的邻居索引。所以，我们需要将`topk_indices`扩展为[batch_size, M, k]吗？是的，因为后续的`gather_nd`需要批次索引。 因此，我们有两种选择： 选择1：保持距离矩阵计算不变（没有批次），然后计算`topk_indices`（形状为[M, k]），然后将其扩展为[batch_size, M, k]（通过tile操作）。 选择2：在计算距离矩阵时，将星座图重复batch_size次，这样距离矩阵就有批次维度。 考虑到星座图在本次训练步骤中是固定的，我们采用选择1，因为更节省内存。 因此，修改步骤： 1. 计算距离矩阵（没有批次维度），得到距离矩阵`distances`（形状[M, M]）。 2. 计算`valid_distances`（形状[M, M]）。 3. 计算`topk_indices`（形状[M, k]）。 4. 将`topk_indices`扩展为[batch_size, M, k]（使用tile）。 但是，注意我们之前已经将距离矩阵计算为[M,M]（没有批次），所以不需要改变。我们只需要调整邻居索引的获取方式。 因此，修改代码： # === 星座点最小距离约束 === points = tf.squeeze(tf.stack([tf.math.real(norm_constellation), tf.math.imag(norm_constellation)], axis=-1)) # 形状: [M, 2] diff = tf.expand_dims(points, 1) - tf.expand_dims(points, 0) # [M, M, 2] distances = tf.norm(diff, axis=-1) # [M, M] # 创建mask（排除自身） mask = tf.eye(tf.shape(distances)[0], dtype=tf.bool) # [M, M] valid_distances = tf.where(mask, tf.ones_like(distances)*1e10, distances) # [M, M] # 计算最小距离（标量） min_distance = tf.reduce_min(valid_distances) distance_penalty = tf.nn.relu(min_distance_threshold - min_distance) * 50.0 # === 新增：局部概率约束（每个点与其最近3个点的概率和≤4/M） === M_float = tf.cast(tf.shape(distances)[0], tf.float32) # 星座点数（浮点数） k = 3 # 最近邻数量 # 1. 找到每个点的最近k个邻居（排除自身）-> 注意没有批次维度 _, topk_indices = tf.math.top_k(-valid_distances, k=k) # 形状: [M, k] # 2. 将topk_indices扩展为批次：形状 [batch_size, M, k] batch_size = tf.shape(p_s)[0] topk_indices_batch = tf.tile(tf.expand_dims(topk_indices, 0), [batch_size, 1, 1]) # 3. 收集邻居概率值 batch_indices = tf.tile(tf.range(batch_size)[:, tf.newaxis, tf.newaxis], [1, tf.shape(topk_indices)[0], k]) # 形状 [batch_size, M, k] gather_indices = tf.stack([batch_indices, topk_indices_batch], axis=-1) # [batch_size, M, k, 2] neighbor_probs = tf.gather_nd(p_s, gather_indices) # [batch_size, M, k] # 4. 计算邻居概率和并检查约束 neighbor_sum = tf.reduce_sum(neighbor_probs, axis=-1) # [batch_size, M] threshold = 4.0 / M_float violations = tf.nn.relu(neighbor_sum - threshold) # [batch_size, M] # 5. 计算惩罚项（平均违反程度×系数） local_prob_penalty = tf.reduce_mean(violations) * 50.0 这样修改后，我们避免了批次维度不一致的问题，并且星座图距离计算只进行一次，然后通过tile扩展到批次。 但是，请注意，在星座图生成过程中，我们使用了`constellation = tf.expand_dims(constellation, 0)`，所以星座图形状为[1, M, 1]（复数），然后归一化后得到`norm_constellation`也是[1, M, 1]（复数）。然后我们计算实部和虚部，堆叠后得到[1, M, 1, 2]，然后squeeze掉第0维和第2维（大小为1的维度），得到[M,2]。所以这里没有批次维度。 因此，我们采用上述修改。 但是，注意：在计算最小距离时，我们使用的是整个星座图的最小距离（标量），所以直接用于损失函数是合理的。 所以，综合以上，我们修改`train_step`函数中局部概率约束部分的代码。 由于原始错误是形状不匹配，我们通过上述修改应该可以解决。 注意：我们修改了局部概率约束的实现，现在使用没有批次维度的距离矩阵，然后将邻居索引扩展到整个批次。 完整修改后的`train_step`函数（仅修改局部概率约束部分）如下： ```python @tf.function def train_step(inp_SNR, noise, GS_flag, PS_flag, eq_flag, epsilon=1e-12, min_distance_threshold=0.5): loss = 0 with tf.GradientTape() as tape: # ... [前面的代码保持不变] ... # === 星座点最小距离约束 === points = tf.squeeze(tf.stack([tf.math.real(norm_constellation), tf.math.imag(norm_constellation)], axis=-1)) # [M, 2] diff = tf.expand_dims(points, 1) - tf.expand_dims(points, 0) # [M, M, 2] distances = tf.norm(diff, axis=-1) # [M, M] mask = tf.eye(tf.shape(distances)[0], dtype=tf.bool) # [M, M] valid_distances = tf.where(mask, tf.ones_like(distances)*1e10, distances) # [M, M] min_distance = tf.reduce_min(valid_distances) distance_penalty = tf.nn.relu(min_distance_threshold - min_distance) * 50.0 # === 新增：局部概率约束（每个点与其最近3个点的概率和≤4/M） === M_int = tf.shape(distances)[0] # 星座点数（整数） M_float = tf.cast(M_int, tf.float32) k = 3 # 最近邻数量 # 1. 找到每个点的最近k个邻居（排除自身） _, topk_indices = tf.math.top_k(-valid_distances, k=k) # [M, k] # 2. 扩展邻居索引到整个批次 batch_size = tf.shape(p_s)[0] topk_indices_batch = tf.tile(tf.expand_dims(topk_indices, 0), [batch_size, 1, 1]) # [batch_size, M, k] # 3. 创建批次索引 batch_indices = tf.tile(tf.range(batch_size)[:, tf.newaxis, tf.newaxis], [1, M_int, k]) # [batch_size, M, k] # 4. 构建收集索引 gather_indices = tf.stack([batch_indices, topk_indices_batch], axis=-1) # [batch_size, M, k, 2] neighbor_probs = tf.gather_nd(p_s, gather_indices) # [batch_size, M, k] # 5. 计算邻居概率和并检查约束 neighbor_sum = tf.reduce_sum(neighbor_probs, axis=-1) # [batch_size, M] threshold = 4.0 / M_float violations = tf.nn.relu(neighbor_sum - threshold) # [batch_size, M] # 6. 计算惩罚项 local_prob_penalty = tf.reduce_mean(violations) * 50.0 # ... [后面的代码保持不变] ... ``` 这样修改后，形状问题应该得到解决。 注意：由于我们修改了`topk_indices`的来源（从无批次维度扩展为有批次维度），并且重新组织了`gather_indices`的构建，所以不会出现形状不匹配的错误。 请尝试运行修改后的代码。</think>这个错误是由于在构建 `gather_indices` 时张量形状不匹配导致的。具体来说，`batch_indices` 是三维张量，而 `topk_indices` 被错误地扩展成了四维张量。让我们修复这个问题： ### 修改后的代码（修复形状不匹配问题） ```python @tf.function def train_step(inp_SNR, noise, GS_flag, PS_flag, eq_flag, epsilon=1e-12, min_distance_threshold=0.5): loss = 0 with tf.GradientTape() as tape: # ... [前面的代码保持不变] ... # === 星座点最小距离约束 === points = tf.squeeze(tf.stack([tf.math.real(norm_constellation), tf.math.imag(norm_constellation)], axis=-1)) diff = tf.expand_dims(points, 1) - tf.expand_dims(points, 0) distances = tf.norm(diff, axis=-1) M_int = tf.shape(distances)[0] mask = tf.eye(M_int, dtype=tf.bool) valid_distances = tf.where(mask, tf.ones_like(distances)*1e10, distances) min_distance = tf.reduce_min(valid_distances) distance_penalty = tf.nn.relu(min_distance_threshold - min_distance) * 50.0 # === 修复：局部概率约束（每个点与其最近3个点的概率和≤4/M） === M_float = tf.cast(M_int, tf.float32) k = 3 # 最近邻数量 # 1. 找到每个点的最近k个邻居（排除自身） _, topk_indices = tf.math.top_k(-valid_distances, k=k) # 形状: [M, k] # 2. 扩展邻居索引到匹配批次维度 batch_size = tf.shape(p_s)[0] topk_indices_exp = tf.expand_dims(topk_indices, 0) # 添加批次维度 topk_indices_batch = tf.tile(topk_indices_exp, [batch_size, 1, 1]) # 形状: [batch_size, M, k] # 3. 创建批次索引 batch_indices = tf.tile(tf.range(batch_size)[:, tf.newaxis, tf.newaxis], [1, M_int, k]) # 形状: [batch_size, M, k] # 4. 构建收集索引（修复形状不匹配） gather_indices = tf.stack([ batch_indices, topk_indices_batch ], axis=-1) # 形状: [batch_size, M, k, 2] # 5. 收集邻居概率值 neighbor_probs = tf.gather_nd(p_s, gather_indices) # 形状: [batch_size, M, k] # 6. 计算邻居概率和并检查约束 neighbor_sum = tf.reduce_sum(neighbor_probs, axis=-1) # 形状: [batch_size, M] threshold = 4.0 / M_float violations = tf.nn.relu(neighbor_sum - threshold) # 7. 计算惩罚项 local_prob_penalty = tf.reduce_mean(violations) * 50.0 # === 概率分布可逆性约束 === # ... [保持不变] ... # === 修改后的损失函数 === loss = (loss_Eq * eq_flag * 0.5 - GMI + loss_NGMI * 100 + distance_penalty + entropy_penalty + prob_floor_penalty + local_prob_penalty) # 梯度计算与更新 variables = (logit_model.trainable_variables * PS_flag + s_model.trainable_variables + Trans_model_bit.trainable_variables * GS_flag + Model_Eq.trainable_variables * eq_flag) gradients = tape.gradient(loss, variables) optimizer.apply_gradients(zip(gradients, variables)) # 返回结果（添加局部概率约束相关指标） return loss, loss_Eq, NGMI, GMI, tf.reduce_mean(entropy_S), p_s, norm_constellation, x, min_distance, local_prob_penalty ``` ### 问题原因与修复解释： 1. **问题原因**： - `topk_indices` 原始形状是 `[M, k]`（M是星座点数，k是邻居数） - 之前的代码尝试使用 `tf.tile(topk_indices[:, :, tf.newaxis, :], [1, 1, k, 1])` 进行扩展，这导致形状变为 `[M, k, k, k]`（四维） - 而 `batch_indices` 是三维张量 `[batch_size, M, k]` - 尝试将三维和四维张量堆叠导致形状不匹配错误 2. **修复方案**： - **步骤1**：保持 `topk_indices` 为原始形状 `[M, k]` - **步骤2**：使用 `tf.expand_dims` 添加批次维度，变为 `[1, M, k]` - **步骤3**：使用 `tf.tile` 沿批次维度复制，得到 `[batch_size, M, k]` - **步骤4**：现在 `batch_indices` 和 `topk_indices_batch` 都是三维 `[batch_size, M, k]` - **步骤5**：使用 `tf.stack` 沿最后一个轴堆叠，得到正确的四维索引张量 `[batch_size, M, k, 2]` 3. **关键改进**： - 确保所有张量在堆叠前具有相同的维度数（rank） - 显式控制张量扩展过程，避免隐式广播导致意外形状 - 保持代码可读性，明确每一步的形状变换 ### 注意事项： 1. 星座图生成部分（`norm_constellation`）应保持为 `[M, 2]` 形状 2. 概率分布 `p_s` 应为 `[batch_size, M]` 形状 3. 邻居数 `k` 应小于星座点数 `M`，否则 `top_k` 会出错 4. 如果批次大小很大，可考虑优化邻居搜索算法以减少计算量