TensorFlow pix2pix模型适应12波段多光谱图像训练指南-Python教程-PHP中文网

TensorFlow pix2pix模型适应12波段多光谱图像训练指南

在图像处理领域，图像到图像的转换任务，如图像风格迁移、超分辨率等，常通过条件生成对抗网络（conditional gans, cgans）实现。tensorflow的pix2pix模型是一个经典的cgan实现，常用于处理三通道rgb图像。然而，当面对多光谱卫星图像这类具有更多波段（如12波段）的数据时，需要对原有的模型架构和数据处理流程进行适应性修改。本文将深入探讨如何将tensorflow的pix2pix模型成功应用于512x512x12维度的多光谱图像数据集，并解决在此过程中可能遇到的张量形状不匹配问题。

1. 适配模型架构以支持多光谱图像

原始的pix2pix模型通常设计用于处理三通道图像。为了适应12波段的多光谱图像，我们需要对生成器（Generator）和判别器（Discriminator）的网络结构进行关键修改，主要是调整输入和输出的通道数。

1.1 生成器（Generator）的修改

生成器的主要任务是将输入图像转换成目标图像。对于12波段图像，其输入和输出都应是12通道。因此，需要调整输入层 Input 的 shape 参数和最终输出层 Conv2DTranspose 的 filters 参数。

import tensorflow as tf
import os
import matplotlib.pyplot as plt

# 假设 downsample 和 upsample 函数已定义，与原pix2pix notebook类似
# downsample 函数通常包含 Conv2D, BatchNorm, LeakyReLU
# upsample 函数通常包含 Conv2DTranspose, BatchNorm, ReLU

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

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

def Generator(output_channels=12): # 增加 output_channels 参数
    input_shape = (512, 512, 12) # 调整输入图像的通道数为12
    inputs = tf.keras.layers.Input(shape=input_shape)

    # 编码器（下采样）层
    down_stack = [
        downsample(64, 4, apply_batchnorm=False),
        downsample(128, 4),
        downsample(256, 4),
        downsample(512, 4),
        downsample(512, 4),
        downsample(512, 4),
        downsample(512, 4),
        downsample(512, 4)
    ]

    # 解码器（上采样）层
    up_stack = [
        upsample(512, 4, apply_dropout=True),
        upsample(512, 4, apply_dropout=True),
        upsample(512, 4, apply_dropout=True),
        upsample(512, 4),
        upsample(256, 4),
        upsample(128, 4),
        upsample(64, 4)
    ]

    initializer = tf.random_normal_initializer(0., 0.02)
    # 最终输出层，通道数应与目标图像的波段数匹配
    last = tf.keras.layers.Conv2DTranspose(output_channels, 4, strides=2,
                                           padding='same',
                                           kernel_initializer=initializer,
                                           activation='tanh')

    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, skips):
        x = up(x)
        if skip is not None:
            x = tf.keras.layers.Concatenate()([x, skip])

    x = last(x)

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

# 实例化生成器，OUTPUT_CHANNELS应设置为12
generator = Generator(output_channels=12)
# generator.summary() # 可用于检查模型结构和参数

登录后复制

1.2 判别器（Discriminator）的修改

判别器负责区分真实图像对（输入图像, 目标真实图像）和生成图像对（输入图像, 生成图像）。因此，判别器的输入也需要调整以适应12波段图像。判别器将输入图像和目标图像（或生成图像）沿通道维度拼接，所以其输入通道数将是 12 + 12 = 24。

def Discriminator():
    initializer = tf.random_normal_initializer(0., 0.02)

    # 输入和目标图像都应有12个通道
    inp = tf.keras.layers.Input(shape=[512, 512, 12], name='input_image')
    tar = tf.keras.layers.Input(shape=[512, 512, 12], name='target_image')

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

    down1 = downsample(64, 4, False)(x)  # (batch_size, 256, 256, 64)
    down2 = downsample(128, 4)(down1)    # (batch_size, 128, 128, 128)
    down3 = downsample(256, 4)(down2)    # (batch_size, 64, 64, 256)

    zero_pad1 = tf.keras.layers.ZeroPadding2D()(down3)  # (batch_size, 66, 66, 256)
    conv = tf.keras.layers.Conv2D(512, 4, strides=1,
                                  kernel_initializer=initializer,
                                  use_bias=False)(zero_pad1)  # (batch_size, 63, 63, 512)

    batchnorm1 = tf.keras.layers.BatchNormalization()(conv)
    leaky_relu = tf.keras.layers.LeakyReLU()(batchnorm1)

    zero_pad2 = tf.keras.layers.ZeroPadding2D()(leaky_relu)  # (batch_size, 65, 65, 512)

    # 最终输出层，通常为1个通道表示真实/伪造的概率
    last = tf.keras.layers.Conv2D(1, 4, strides=1,
                                  kernel_initializer=initializer)(zero_pad2)  # (batch_size, 62, 62, 1)

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

# 实例化判别器
discriminator = Discriminator()
# discriminator.summary() # 可用于检查模型结构和参数

登录后复制

2. 处理多光谱图像数据加载与批次维度

多光谱图像通常以多种格式存储，例如将每3个波段存储为一个RGB图像（16位PNG），然后组合成12波段图像。在数据加载函数 load() 中，需要确保将所有相关波段正确读取并组合成 (512, 512, 12) 的张量。

2.1 关键的批次维度问题

在TensorFlow中，模型期望的输入张量通常包含一个批次维度（batch dimension），即形状为 (batch_size, height, width, channels)。当使用 tf.data.Dataset.take(N) 迭代数据集时，如果 N 是1或较小的数字，take() 方法会返回单个样本，而不是一个批次。这意味着 input_image 和 target_image 的形状将是 (height, width, channels)，缺少批次维度，从而导致模型输入形状不匹配错误，例如 ValueError: Input 0 of layer "model_1" is incompatible with the layer: expected shape=(None, 512, 512, 12), found shape=(512, 512, 12)。

解决此问题有两种主要方法：

手动添加批次维度： 在将单个样本传递给模型之前，使用 tf.expand_dims(tensor, 0) 手动添加批次维度。
使用 Dataset.batch() 方法： 在创建数据集时，显式地使用 dataset.batch(batch_size) 方法对数据进行批处理。这是更推荐的做法，因为它与TensorFlow的数据管道设计更为契合，并能提高训练效率。

推荐的数据集批处理方式：

# 假设 train_raw_dataset 已经包含了原始图像路径或数据
# 定义一个加载和预处理图像的函数，该函数应返回 (input_image, target_image)
# 例如，load_image 函数会读取12波段图像并进行归一化
# def load_image(image_path):
#    # ... 读取12波段图像的逻辑 ...
#    input_image = ... # (512, 512, 12)
#    target_image = ... # (512, 512, 12)
#    return input_image, target_image

# train_dataset = train_raw_dataset.map(load_image, num_parallel_calls=tf.data.AUTOTUNE)
# test_dataset = test_raw_dataset.map(load_image, num_parallel_calls=tf.data.AUTOTUNE)

# 使用 .batch() 方法添加批次维度
BATCH_SIZE = 1 # 或者更大的批次大小，取决于GPU内存
train_dataset = train_dataset.batch(BATCH_SIZE).prefetch(tf.data.AUTOTUNE)
test_dataset = test_dataset.batch(BATCH_SIZE).prefetch(tf.data.AUTOTUNE)

# 训练循环示例 (fit函数内部)
# for step, (input_image, target) in train_dataset.enumerate():
#    # input_image.shape 现在会是 (BATCH_SIZE, 512, 512, 12)
#    # ... 训练步骤 ...

登录后复制

2.2 调整 generate_images 函数

generate_images 函数用于可视化模型的生成结果。由于模型现在处理的是12波段图像，而标准显示器通常只支持3个通道（RGB），我们需要修改此函数以分批次显示3个波段，或者选择特定的3个波段进行可视化。同时，为了与模型输出兼容，输入到此函数中的 test_input 和 tar 也应具有批次维度。

可图大模型

可图大模型（Kolors）是快手大模型团队自研打造的文生图AI大模型

查看详情

def generate_images(model, test_input, tar, save_dir='generated_images'):
    # 确保 test_input 和 tar 具有批次维度
    if len(test_input.shape) != 4:
        test_input = tf.expand_dims(test_input, 0)
    if len(tar.shape) != 4:
        tar = tf.expand_dims(tar, 0)

    prediction = model(test_input, training=True)

    num_bands = 12
    # 每次显示3个波段，例如：0-2, 3-5, 6-8, 9-11
    for i in range(0, num_bands, 3):
        # 确定当前要显示的波段索引
        bands = [i, i + 1, i + 2]
        # 处理最后一个分组可能不足3个波段的情况
        bands = [b for b in bands if b < num_bands]
        if not bands:
            continue

        plt.figure(figsize=(15, 5))

        display_list = [test_input[0], tar[0], prediction[0]]
        title = ['Input Image', 'Ground Truth', 'Predicted Image']

        for j in range(3):
            plt.subplot(1, 3, j + 1)
            plt.title(title[j])

            # 选择并堆叠指定波段进行可视化
            # 确保即使 bands 不足3个，也能正确堆叠
            image_display = tf.stack([display_list[j][..., band] for band in bands], axis=-1)

            # 如果选择的波段不足3个，可以填充或调整显示方式
            if image_display.shape[-1] < 3:
                # 简单填充，例如复制最后一个波段，或者根据需求调整
                if image_display.shape[-1] == 1:
                    image_display = tf.concat([image_display, image_display, image_display], axis=-1)
                elif image_display.shape[-1] == 2:
                    image_display = tf.concat([image_display, image_display[..., -1:]], axis=-1)

            # 将图像数据重新缩放到 [0, 1] 范围以便显示
            image_display = (image_display + 1) / 2

            plt.imshow(image_display)
            plt.axis('off')

        os.makedirs(save_dir, exist_ok=True)
        # 保存图像时，明确指出显示的是哪些波段
        band_str = "_".join(map(str, bands))
        plt.savefig(os.path.join(save_dir, f'generated_image_bands_{band_str}.png'))
        plt.close()

# 示例用法
# for example_input, example_target in test_dataset.take(1):
#     generate_images(generator, example_input, example_target)

登录后复制

注意事项：

在 generate_images 函数中，test_input[0]、tar[0] 和 prediction[0] 假定 test_input、tar 和 prediction 都具有批次维度。如果 test_dataset.take(1) 返回的是无批次维度的单个样本，则需要像代码中所示，在传递给 model 之前手动添加批次维度，或者在 test_dataset 上使用 .batch(1)。
image_display = (image_display + 1) / 2 这一行假设模型输出的像素值范围是 [-1, 1]，这是GANs中常见的激活函数 tanh 的输出范围。如果模型输出范围不同，需要相应调整归一化方法。

3. 调试与错误排查

在训练过程中，可能会遇到 InvalidArgumentError 或 tf.function retracing 警告。

3.1 InvalidArgumentError 错误排查

这类错误通常与张量形状不匹配有关。除了上述的批次维度问题，还需要检查：

数据加载函数 load()： 确保它始终输出 (height, width, channels) 形状的张量，且 channels 数量与模型期望的输入通道数一致（例如12）。
数据集的一致性： 检查所有输入图像是否都具有相同的 (512, 512, 12) 形状。不一致的输入形状会导致 tf.function retracing 警告，甚至更严重的运行时错误。
模型内部层： 仔细检查 Generator 和 Discriminator 中所有层的输入和输出形状是否按预期传递。generator.summary() 和 discriminator.summary() 是非常有用的工具，可以打印出每一层的形状信息。

使用虚拟数据进行快速测试： 为了快速验证模型架构本身是否正确，可以使用随机生成的虚拟数据来创建一个简单的 tf.data.Dataset。这有助于排除数据加载或预处理阶段的问题。

# 生成100个虚拟的 512x512x12 图像数据
x_dummy = tf.random.uniform((100, 512, 512, 12), minval=-1, maxval=1, dtype=tf.float32)
# 创建一个虚拟数据集，其中输入和目标图像相同
dummy_train_dataset = tf.data.Dataset.from_tensor_slices((x_dummy, x_dummy))
# 对虚拟数据集进行批处理
dummy_train_dataset = dummy_train_dataset.batch(BATCH_SIZE).prefetch(tf.data.AUTOTUNE)

# 使用虚拟数据集进行训练测试
# fit(dummy_train_dataset, test_dataset, steps=...)

登录后复制

如果使用虚拟数据训练没有报错，那么问题很可能出在您的真实数据加载、预处理或数据集构建过程中。

3.2 tf.function Retracing 警告

WARNING:tensorflow:X out of the last X calls to <function ...> triggered tf.function retracing. 这个警告表明 tf.function 被多次重新追踪（retrace），这会降低训练效率。常见原因包括：

传入不同形状的张量： 最常见的原因是每次调用 tf.function 时，传入的张量形状不一致。这通常发生在数据集中的图像尺寸不统一，或者如前所述，批次维度有时存在有时缺失。确保通过 dataset.batch() 方法获得的数据批次始终具有相同的形状。
传入Python对象而非张量： 避免在 tf.function 内部使用Python列表、字典等作为模型输入，应将其转换为TensorFlow张量。
在循环中定义 tf.function： 确保 tf.function 装饰的函数在循环外部定义。

检查数据管道，确保每个批次的 input_image 和 target_image 都具有一致的 (BATCH_SIZE, 512, 512, 12) 形状是解决此警告的关键。

4. 总结

将TensorFlow pix2pix模型应用于12波段多光谱图像需要对模型架构和数据处理流程进行细致的调整。核心步骤包括：

修改生成器和判别器： 调整 Input 层的通道数和生成器最终 Conv2DTranspose 层的 filters 参数，使其与12波段输入/输出图像匹配。
数据管道优化： 确保数据加载函数能够正确读取并组合12波段图像。最重要的是，通过 tf.data.Dataset.batch() 方法对数据进行批处理，以确保模型始终接收到带有批次维度的张量 (BATCH_SIZE, height, width, channels)，从而避免形状不匹配错误。
可视化函数适配： 更新 generate_images 函数以正确处理和显示多波段图像，通常通过分批次显示3个波段来实现。
调试策略： 利用 model.summary() 检查模型结构，并使用虚拟数据快速验证模型架构的正确性。对 tf.function retracing 警告保持警惕，并检查数据管道中的形状一致性。

通过遵循这些专业指南，您可以有效地将pix2pix模型应用于多光谱图像数据，实现高质量的图像到图像转换任务。

以上就是TensorFlow pix2pix模型适应12波段多光谱图像训练指南的详细内容，更多请关注php中文网其它相关文章！