Keras多GPU数据并行模式下提取嵌入层权重的问题

这个问题我之前帮不少开发者解决过——多GPU数据并行下模型结构变形，本质是Keras的分布式机制会对原模型做封装/复制，导致你原本的单GPU逻辑和多GPU环境不兼容。下面给你拆解原因和解决方案：

问题根源 #

当你使用Keras官方的多GPU数据并行方案（比如tf.distribute.MirroredStrategy或者旧版的tf.keras.utils.multi_gpu_model）时，系统会自动在每个GPU上创建原模型的副本，然后把输入数据拆分到各个GPU上并行计算，最后汇总梯度更新权重。这个过程中，原模型会被包装成一个分布式专用的模型结构，所以你看到的模型摘要会和单GPU模式下不一样；更关键的是，你直接提取的权重可能是分布式模型副本中的权重，而非你最初定义的那个全连接层的权重。

解决方案 #

方案1：分离训练与嵌入计算流程（最推荐） #

训练阶段用多GPU加速，训练完成后把权重转移回单GPU的基础模型，再用基础模型处理嵌入权重提取和点积计算：

import tensorflow as tf

# 1. 先定义单GPU的基础模型（包含你的嵌入全连接层）
base_model = tf.keras.Sequential([
    tf.keras.layers.Dense(256, activation='relu', input_shape=(input_dim,)),
    tf.keras.layers.Dense(embedding_dim, name='image_embedding')  # 你的嵌入层
])

# 2. 配置多GPU训练策略
strategy = tf.distribute.MirroredStrategy()
with strategy.scope():
    # 复制基础模型到多GPU环境
    multi_gpu_model = tf.keras.models.clone_model(base_model)
    multi_gpu_model.compile(optimizer='adam', loss='categorical_crossentropy')

# 3. 用多GPU模型训练
multi_gpu_model.fit(train_data, train_labels, epochs=20)

# 4. 训练完成后，把多GPU模型的权重同步回基础模型
base_model.set_weights(multi_gpu_model.get_weights())

# 5. 现在用基础模型提取嵌入权重，执行点积运算
embedding_weights = base_model.get_layer('image_embedding').get_weights()[0]
# 假设input_vectors是你的输入向量组，shape为[batch_size, input_dim]
dot_product = tf.matmul(input_vectors, embedding_weights, transpose_b=True)

方案2：将点积逻辑整合为模型的一部分 #

如果你需要在多GPU环境下直接完成点积计算，可以把点积操作封装成自定义层，整合到模型中，这样分布式策略会自动处理多GPU同步，模型结构也会正常显示：

import tensorflow as tf

class EmbeddingDotProduct(tf.keras.layers.Layer):
    def __init__(self, embedding_layer_name):
        super().__init__()
        self.embedding_layer_name = embedding_layer_name

    def build(self, input_shape):
        # 获取嵌入层的权重（会自动同步多GPU的权重）
        self.embedding_weights = self.get_layer(self.embedding_layer_name).get_weights()[0]
        super().build(input_shape)

    def call(self, inputs):
        # inputs是你的输入向量组
        return tf.matmul(inputs, self.embedding_weights, transpose_b=True)

# 在多GPU策略下定义完整模型
strategy = tf.distribute.MirroredStrategy()
with strategy.scope():
    # 主输入（用于训练嵌入层）
    main_input = tf.keras.layers.Input(shape=(input_dim,))
    x = tf.keras.layers.Dense(256, activation='relu')(main_input)
    embedding_layer = tf.keras.layers.Dense(embedding_dim, name='image_embedding')(x)
    
    # 点积计算的输入向量组
    vector_input = tf.keras.layers.Input(shape=(input_dim,))
    dot_product_output = EmbeddingDotProduct('image_embedding')(vector_input)
    
    # 构建完整模型
    full_model = tf.keras.Model(inputs=[main_input, vector_input], outputs=[embedding_layer, dot_product_output])
    full_model.compile(optimizer='adam', loss=['categorical_crossentropy', None])

# 训练时同时传入主输入和向量输入（或按需调整训练逻辑）
full_model.fit([train_main_data, train_vector_data], [train_labels, None], epochs=20)

注意事项 #

• 不要直接从多GPU包装后的模型中提取层权重，因为分布式模型的层是被封装的，结构显示会混乱，且权重提取逻辑可能不符合预期。
• 多GPU环境下，所有需要跨GPU同步的操作都要放在strategy.scope()上下文管理器中，或者通过Keras的层/模型API来实现，避免手动操作权重导致的同步问题。

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