自定义TensorFlow Embedding层GPU无法高效运行问题排查

架构老林

2026-8-20

问题：TensorFlow模型GPU利用率极低的原因与解决方法

模型实现

EnhancedEmbedding层

class EnhancedEmbedding(tf.keras.layers.Embedding):
    def __init__(self, input_dim, output_dim, embeddings_initializer='uniform', embeddings_regularizer=None, activity_regularizer=None, embeddings_constraint=None, mask_zero=False, input_length=None, **kwargs):
        super().__init__(input_dim, output_dim, embeddings_initializer, embeddings_regularizer, activity_regularizer, embeddings_constraint, mask_zero, input_length, **kwargs)
        self.five=tf.constant(0.5)
        self.zero=tf.constant(0)
    
    def embedding(self,inputs):
        return super().call(inputs)
    
    def map_2(self,tokens):
        identifier=self.embedding(tokens[0])
        cur_word=self.embedding(tokens[2])
        return tf.cond(tf.equal(tf.shape(tokens[0]),self.zero),
                lambda: tf.squeeze(cur_word),
                lambda: tf.squeeze(tf.reduce_mean(identifier)*self.five+cur_word*self.five))


    def map_1(self,inputs):
        return tf.map_fn(fn=lambda x :self.map_2(x),elems=inputs,dtype=tf.float32)

    def call(self, inputs):
        final_embeddings=tf.map_fn(fn=lambda x:self.map_1(x),elems=inputs,dtype=tf.float32)

        return final_embeddings

EnhancedModel模型

class EnhancedModel(Model):
    def __init__(self,  embedding_dim, hidden_dim, vocab_size, label_size,seq_len, pretrained_weight):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.vocab_size = vocab_size
        self.embedding_dim = embedding_dim
        self.label_size = label_size
        self.activation = tf.keras.activations.tanh
        self.num_layers = 1

        self.embedding=EnhancedEmbedding(vocab_size,embedding_dim,embeddings_initializer=keras.initializers.Constant(pretrained_weight))
        self.encoder = Bidirectional(LSTM(hidden_dim, return_sequences=True,input_shape=(seq_len,embedding_dim)))
        self.pool=MaxPool1D(hidden_dim*2)
        self.decoder = Dense(self.label_size)

    def call(self, inputs):
        embeddings = self.embedding(inputs)
        embeddings=tf.reshape(embeddings,[-1,400,32])
        lstm_out = self.encoder(embeddings)
        lstm_out = tf.transpose(lstm_out, perm=[0,2,1])
        pool_out=self.pool(lstm_out)
        out = tf.squeeze(pool_out,[1])
        out = self.decoder(out)
        return out

环境配置与核心问题

Tesla-V100环境：TensorFlow 2.3.0、cudatoolkit 10.1、cudnn 7.6.5，GPU可正常调用（tf.test.is_gpu_available()返回True）
问题现象：输入为形状(batch_size,400,3,None)的RaggedTensor，模型功能正常，但GPU利用率仅约5%，CPU利用率达100%

后续测试结果

在笔记本环境（i7-12700H、TensorFlow 2.9.1、cuda 11.6、cudnn 8.4.1）、batch_size=64的条件下：

关闭Eager模式（tf.config.run_functions_eagerly(False)）：单批次训练耗时约3分钟，大部分时间CPU利用率20%、GPU利用率17%，仅少数时间GPU利用率超50%
开启Eager模式（tf.config.run_functions_eagerly(True)）：GPU利用率全程维持在80%以上，CPU利用率仅9%
移除map_2中的tf.cond后，Tesla-V100上单批次训练耗时仍约100秒，无明显改善

原因分析

嵌套tf.map_fn的低效性：tf.map_fn在Graph模式下会生成大量细粒度运算节点，无法利用GPU的并行计算优势，反而引发CPU-GPU数据传输瓶颈。尤其是处理RaggedTensor时，动态形状的逐元素遍历会导致计算图优化失效，大部分操作 fallback到CPU执行。
RaggedTensor的Graph模式优化不足：TensorFlow 2.3对RaggedTensor的Graph模式支持不完善，即使升级到2.9版本，嵌套map_fn结合RaggedTensor的计算图优化依然不佳。而Eager模式下，TensorFlow直接在GPU执行张量操作，避免了Graph模式的节点调度和静态形状推断开销。
tf.cond的影响有限：虽然tf.cond会增加计算图分支复杂度，但移除后性能无明显提升，说明核心瓶颈并非条件分支，而是嵌套遍历的低效性。

解决方法

替换嵌套tf.map_fn为向量化操作：利用TensorFlow的批量API处理RaggedTensor，避免逐元素遍历。示例修改EnhancedEmbedding的call方法：

def call(self, inputs):
    # 提取identifier和cur_word对应的RaggedTensor
    identifier = inputs[:, :, 0]  # 形状: (batch_size, 400, None)
    cur_word = inputs[:, :, 2]    # 形状: (batch_size, 400, 1)（假设cur_word为单个token）
    
    # 批量计算embedding
    identifier_emb = super().call(identifier)  # 形状: (batch_size, 400, None, embedding_dim)
    cur_word_emb = super().call(cur_word)      # 形状: (batch_size, 400, 1, embedding_dim)
    
    # 对identifier的可变维度求均值，对cur_word挤压冗余维度
    identifier_mean = tf.reduce_mean(identifier_emb, axis=2)  # 形状: (batch_size, 400, embedding_dim)
    cur_word_squeezed = tf.squeeze(cur_word_emb, axis=2)       # 形状: (batch_size, 400, embedding_dim)
    
    # 加权合并得到最终embedding
    final_embeddings = 0.5 * identifier_mean + 0.5 * cur_word_squeezed
    return final_embeddings

注：需根据cur_word的实际形状调整squeeze的axis参数，确保维度匹配。

升级TensorFlow版本：使用TensorFlow 2.10+版本，该版本对RaggedTensor的Graph模式支持和向量化操作优化更完善，能有效减少CPU fallback情况。
优先使用批量操作API：自定义层中避免使用循环或逐元素遍历，尽量利用TensorFlow的内置批量运算（如tf.reduce_mean、tf.concat等），让计算图生成高效的GPU执行算子。
必要时保留Eager模式训练：若向量化改造难度大，在TensorFlow 2.9+版本中，Eager模式的GPU性能已大幅提升，可保持tf.config.run_functions_eagerly(True)进行训练，实际效率优于优化不佳的Graph模式。

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

自定义TensorFlow Embedding层GPU无法高效运行问题排查

模型实现 #

EnhancedEmbedding层 #

EnhancedModel模型 #

环境配置与核心问题 #

后续测试结果 #

原因分析 #

解决方法 #