TensorFlow中神经网络离散颜色选择转RGB图像及报错修复

我正在学习TensorFlow，计划实现一个抖动神经网络（dithering neural network）实验：为每个像素提供4种颜色选项，让网络学习为每个像素挑选最优颜色以还原输入图像。需要将网络输出的形状为(width*height, 4)的张量映射为RGB图像，与模糊后的参考图像对比。

已有可实现该功能的NumPy版本代码，但该代码无法用于强化学习场景，因为TensorFlow无法基于它计算网络梯度。尝试编写TensorFlow版本的转换代码时，出现索引越界错误：IndexError: index 4 is out of bounds for axis 1 with size 4。

NumPy版本转换代码 #

# NumPy版本转换代码
def convert_to_image(output_matrix, color_choices):
    """Maps network outputs to RGB colors using predefined choices."""
    # Converts (height*width, 4-floats) to (height*width) array of 0-3 indices
    chosen_colors = np.argmax(output_matrix, axis=1)
    # Uses color_choices to match each indice to a pixel color
    mapped_image = color_choices[np.arange(color_choices.shape[0]), chosen_colors] * 255

    mapped_image = mapped_image.reshape((INPUT_SHAPE[0], INPUT_SHAPE[1], 3))  # Reshape to (height, width, 3)
    return mapped_image

TensorFlow版本尝试代码 #

# TensorFlow版本尝试代码
def convert_to_image_tf(output_matrices, color_choices_batch):
    """Maps network outputs to RGB colors using predefined choices in TensorFlow."""
    batch_size = output_matrices.shape[0]
    print("output_matrices dimensions: " + str(output_matrices.shape))
    print("color_choices_batch dimensions: " + str(color_choices_batch.shape))
    chosen_colors = tf.argmax(output_matrices, axis=2, output_type=tf.int32)
    print("chosen_colors shape: " + str(chosen_colors.shape))

    batch_size, num_pixels, num_choices_per_px, num_colors = color_choices_batch.shape  # Extract dimensions

    # Use batch-wise advanced indexing
    mapped_image = color_choices[np.arange(batch_size)[:, None], np.arange(num_pixels), chosen_colors] * 255

    print("mapped_image shape: " + str(mapped_image.shape))
    mapped_image = tf.reshape(mapped_image, (INPUT_SHAPE[0], INPUT_SHAPE[1], 3))  # Reshape to (height, width, 3)

    return mapped_image

报错信息 #

output_matrices dimensions: (32, 1024, 4)
color_choices_batch dimensions: (32, 1024, 4, 3)
chosen_colors shape: (32, 1024)
Traceback (most recent call last):
  File "xxx\python_tests\dither_test_1.py", line 229, in <module>
    loss = train_step(color_decider, image_batch, color_choices_batch)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "xxx\python_tests\dither_test_1.py", line 209, in train_step
    generated_images = convert_to_image_tf(raw_outputs, color_choices_batch)
                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "xxx\python_tests\dither_test_1.py", line 182, in convert_to_image_tf
    mapped_image = color_choices[np.arange(batch_size)[:, None], np.arange(num_pixels), chosen_colors] * 255
                   ~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
IndexError: index 4 is out of bounds for axis 1 with size 4

问题分析与修复 #

1. 变量名错误：代码中误用了未定义的color_choices变量，实际应该使用传入的参数color_choices_batch，这是引发索引错误的直接原因。
2. 索引方式不兼容：不能直接混用NumPy数组和TensorFlow张量做高级索引，TensorFlow有原生的批量索引API。
3. 维度匹配问题：原代码忽略了batch维度的处理，reshape逻辑会导致形状不匹配。

修复后的TensorFlow代码 #

def convert_to_image_tf(output_matrices, color_choices_batch):
    """Maps network outputs to RGB colors using predefined choices in TensorFlow."""
    batch_size = tf.shape(output_matrices)[0]
    num_pixels = tf.shape(output_matrices)[1]
    
    # 获取每个像素选择的颜色索引，形状(32,1024)
    chosen_colors = tf.argmax(output_matrices, axis=2, output_type=tf.int32)
    
    # 构造batch维度索引，形状(32,1024,1)
    batch_indices = tf.tile(tf.reshape(tf.range(batch_size), (-1,1,1)), (1, num_pixels, 1))
    # 构造像素维度索引，形状(32,1024,1)
    pixel_indices = tf.tile(tf.reshape(tf.range(num_pixels), (1,-1,1)), (batch_size, 1, 1))
    # 合并为三维索引，匹配color_choices_batch的(batch, pixel, choice)维度
    gather_indices = tf.concat([batch_indices, pixel_indices, tf.expand_dims(chosen_colors, axis=-1)], axis=-1)
    
    # 按索引提取对应颜色，形状(32,1024,3)
    mapped_image = tf.gather_nd(color_choices_batch, gather_indices) * 255
    
    # 调整为图像形状，保留batch维度
    mapped_image = tf.reshape(mapped_image, (batch_size, INPUT_SHAPE[0], INPUT_SHAPE[1], 3))
    return mapped_image

关键修复点 #

• 修正变量名，使用传入的color_choices_batch替代未定义的color_choices
• 使用TensorFlow原生tf.gather_nd实现批量索引，避免跨框架类型混用
• 构造正确的三维索引张量，与输入张量维度完全对齐
• 保留batch维度的reshape逻辑，适配批量训练场景

内容的提问来源于stack exchange，提问作者Tomáš Zato