图像补丁形状适配问题：将(81,256)调整为(1,81,256)及批量支持

问题核心在于处理图像补丁时丢失了批次维度，导致最终输出无法区分不同图像对应的补丁组。以下是具体的问题定位和修复步骤：

问题定位 #

模型代码中这一步操作直接将批次维度和补丁维度合并，是丢失维度信息的关键原因：

patches = tf.reshape(patches, [-1, 32, 32, 3])

原本Patches层输出的形状是(批次大小, 补丁数, 15552)，经过上述reshape后变成(批次大小*补丁数, 32,32,3)，后续传入Xception处理时，无法区分哪些补丁属于同一图像，最终输出丢失了批次维度。

修复方案 #

需要保留(批次大小, 补丁数, 32,32,3)的层级结构，使用TimeDistributed层将Xception模型应用到每个补丁上，以此保留批次与补丁的对应关系。

修复后的完整模型代码 #

xception = keras.applications.Xception(
    include_top=False, weights="imagenet", pooling="avg"
)

for layer in xception.layers:
    layer.trainable = trainable

# 用TimeDistributed包裹Xception，实现对每个补丁的独立处理
xception_distributed = tf.keras.layers.TimeDistributed(xception)

inputs = layers.Input(shape=(300, 300, 3), name="image_input")
patch_size = 72

print("inputs shape is ", inputs.shape)

xception_input = tf.keras.applications.xception.preprocess_input(inputs)
print("xception shape is", xception_input.shape)
patches = Patches(patch_size)(xception_input)

print("patches shape is", patches.shape)
# 显式保留批次和补丁维度，避免合并
batch_size = tf.shape(patches)[0]
num_patches = tf.shape(patches)[1]
patches = tf.reshape(patches, [batch_size, num_patches, 32, 32, 3])

print("patches shape is", patches.shape)

# 对每个批次内的所有补丁应用Xception
embeddings = xception_distributed(patches)

output = embeddings(
    embeddings, projection_layers, dims, dropout_rate
)

print("output shape is ", output.shape)
return keras.Model(inputs, output)

关键修改说明 #

1. TimeDistributed层：包裹Xception后，会自动对输入的(批次大小, 补丁数, 32,32,3)中的每个(32,32,3)补丁单独运行Xception，输出形状变为(批次大小, 补丁数, 2048)（Xception全局平均池化后的维度）。
2. 维度保留的reshape：不再用-1合并批次与补丁维度，而是基于patches的前两个维度显式reshape，确保层级结构清晰。

修改完成后，单张图像输入时输出形状为(1,81,256)，20张图像的批次输入时输出形状为(20,81,256)，完全符合需求。

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