基于TensorFlow pix2pix训练12波段卫星图像模型的张量错误排查

核心问题定位 #

原Pix2Pix针对256×256×3的RGB图像设计，下采样/上采样的维度计算、通道数匹配逻辑都是基于该规格。当切换到512×512×12的多光谱图像时，若仅简单修改通道数和尺寸数值，容易出现下采样/上采样层数不匹配、张量拼接形状冲突、硬编码尺寸残留等问题，进而触发切片越界、形状不匹配的错误。

具体修复步骤 #

1. 强制保证数据加载后的形状一致性

样本形状不一致是随机触发错误的常见原因，需在数据管道中添加严格的形状校验：

def validate_sample_shape(input_img, target_img):
    # 强制设置固定形状，避免读取错误导致的尺寸偏差
    input_img.set_shape((512, 512, 12))
    target_img.set_shape((512, 512, 12))
    return input_img, target_img

# 应用到训练和测试数据集
train_dataset = train_dataset.map(validate_sample_shape)
test_dataset = test_dataset.map(validate_sample_shape)

# 验证数据集元素形状
print("训练集元素形状:", train_dataset.element_spec)
# 预期输出: (TensorSpec(shape=(512, 512, 12), dtype=tf.float32, name=None), TensorSpec(shape=(512, 512, 12), dtype=tf.float32, name=None))

同时检查16位转浮点的归一化过程，确保转换后张量形状未发生变化。

2. 修正生成器(Generator)的下采样/上采样层数

512是2的9次幂，需将原8层下采样/上采样调整为9层，保证特征图维度能完全还原：

def downsample(filters, size, apply_batchnorm=True):
    initializer = tf.random_normal_initializer(0., 0.02)
    seq = tf.keras.Sequential()
    seq.add(tf.keras.layers.Conv2D(filters, size, strides=2, padding='same',
                                  kernel_initializer=initializer, use_bias=False))
    if apply_batchnorm:
        seq.add(tf.keras.layers.BatchNormalization())
    seq.add(tf.keras.layers.LeakyReLU())
    return seq

def upsample(filters, size, apply_dropout=False):
    initializer = tf.random_normal_initializer(0., 0.02)
    seq = tf.keras.Sequential()
    seq.add(tf.keras.layers.Conv2DTranspose(filters, size, strides=2,
                                           padding='same',
                                           kernel_initializer=initializer,
                                           use_bias=False))
    seq.add(tf.keras.layers.BatchNormalization())
    if apply_dropout:
        seq.add(tf.keras.layers.Dropout(0.5))
    seq.add(tf.keras.layers.ReLU())
    return seq

# 针对512×512的下采样栈（9层，最终得到1×1特征图）
down_stack = [
    downsample(64, 4, apply_batchnorm=False),  # (bs, 256, 256, 64)
    downsample(128, 4),  # (bs, 128, 128, 128)
    downsample(256, 4),  # (bs, 64, 64, 256)
    downsample(512, 4),  # (bs, 32, 32, 512)
    downsample(512, 4),  # (bs, 16, 16, 512)
    downsample(512, 4),  # (bs, 8, 8, 512)
    downsample(512, 4),  # (bs, 4, 4, 512)
    downsample(512, 4),  # (bs, 2, 2, 512)
    downsample(512, 4),  # (bs, 1, 1, 512)
]

# 对应的上采样栈（9层，最终输出512×512×12）
up_stack = [
    upsample(512, 4, apply_dropout=True),  # (bs, 2, 2, 1024)
    upsample(512, 4, apply_dropout=True),  # (bs, 4, 4, 1024)
    upsample(512, 4, apply_dropout=True),  # (bs, 8, 8, 1024)
    upsample(512, 4),  # (bs, 16, 16, 1024)
    upsample(512, 4),  # (bs, 32, 32, 1024)
    upsample(256, 4),  # (bs, 64, 64, 512)
    upsample(128, 4),  # (bs, 128, 128, 256)
    upsample(64, 4),  # (bs, 256, 256, 128)
    upsample(12, 4),  # (bs, 512, 512, 12)
]

# 生成器前向逻辑，确保跳连形状匹配
def Generator():
    inputs = tf.keras.layers.Input(shape=[512, 512, 12])
    x = inputs

    # 下采样收集跳连特征
    skips = []
    for down in down_stack:
        x = down(x)
        skips.append(x)
    skips = reversed(skips[:-1])  # 跳过最后一层下采样的特征

    # 上采样拼接跳连特征
    for up, skip in zip(up_stack[:-1], skips):
        x = up(x)
        x = tf.keras.layers.Concatenate()([x, skip])

    # 最后一层上采样输出目标图像
    last = up_stack[-1](x)
    outputs = tf.keras.layers.Activation('tanh')(last)

    return tf.keras.Model(inputs=inputs, outputs=outputs)

3. 修正判别器(Discriminator)的输入通道与维度

判别器需要拼接输入图像和目标/生成图像，因此输入通道数为12+12=24，同时调整下采样层数适配512×512：

def Discriminator():
    initializer = tf.random_normal_initializer(0., 0.02)
    inp = tf.keras.layers.Input(shape=[512, 512, 12], name='input_image')
    tar = tf.keras.layers.Input(shape=[512, 512, 12], name='target_image')

    # 拼接输入和目标图像，通道数24
    x = tf.keras.layers.concatenate([inp, tar])  # (bs, 512, 512, 24)

    # 下采样到1×1特征图
    down1 = downsample(64, 4, apply_batchnorm=False)(x)  # (bs, 256, 256, 64)
    down2 = downsample(128, 4)(down1)  # (bs, 128, 128, 128)
    down3 = downsample(256, 4)(down2)  # (bs, 64, 64, 256)
    down4 = downsample(512, 4)(down3)  # (bs, 32, 32, 512)
    down5 = downsample(512, 4)(down4)  # (bs, 16, 16, 512)
    down6 = downsample(512, 4)(down5)  # (bs, 8, 8, 512)
    down7 = downsample(512, 4)(down6)  # (bs, 4, 4, 512)
    down8 = downsample(512, 4)(down7)  # (bs, 2, 2, 512)

    # 输出真假判断的单通道特征图
    last = tf.keras.layers.Conv2D(1, 4, strides=1,
                                  kernel_initializer=initializer)(down8)  # (bs, 1, 1, 1)

    return tf.keras.Model(inputs=[inp, tar], outputs=last)

4. 移除代码中硬编码的256尺寸

检查generate_images()、数据增强等自定义函数，将所有硬编码的256替换为动态获取的张量形状，比如：

def generate_images(model, test_input, tar):
    prediction = model(test_input, training=True)
    # 动态获取图像尺寸，避免硬编码
    img_height = tf.shape(test_input)[1]
    img_width = tf.shape(test_input)[2]

    # 多波段可视化示例：取前3波段显示
    plt.figure(figsize=(18, 6))
    display_list = [test_input[0,...,0:3], tar[0,...,0:3], prediction[0,...,0:3]]
    titles = ['Input (Bands 1-3)', 'Ground Truth (Bands 1-3)', 'Prediction (Bands 1-3)']

    for i in range(3):
        plt.subplot(1, 3, i+1)
        plt.title(titles[i])
        # 归一化到[0,1]范围用于显示
        plt.imshow(tf.keras.utils.normalize(display_list[i], axis=0))
        plt.axis('off')
    plt.show()

5. 预校验模型输出形状

训练前手动传入测试张量，验证生成器和判别器的输出形状是否符合预期：

# 生成测试张量
test_input = tf.random.normal([1, 512, 512, 12])
test_target = tf.random.normal([1, 512, 512, 12])

# 验证生成器
generator = Generator()
gen_output = generator(test_input, training=False)
print(f"生成器输出形状: {gen_output.shape}")  # 预期: (1, 512, 512, 12)

# 验证判别器
discriminator = Discriminator()
disc_output = discriminator([test_input, gen_output], training=False)
print(f"判别器输出形状: {disc_output.shape}")  # 预期: (1, 1, 1, 1)

若形状不符合预期，先修复模型结构再启动训练。

内容的提问来源于stack exchange，提问作者Naser.Sadeghi