如何在TensorFlow2中高效编码多类别特征？解决Embedding层存储过大问题

问题背景 #

在TensorFlow 2.2.0中实现多类别特征的Embedding模块，输入形状为batch_size * num_features，每一列是整数编码的类别特征（如某特征有10个类别，则取值为0-9）。最初采用创建layers.Embedding()列表的方式为每个特征生成Embedding，但保存模型时出现异常：当num_features=3000、每个特征10个类别、Embedding维度为3时，理论参数总量为90000（3000103），对应存储约0.34M，但model.save()生成的saved_model.pb文件达59M。尝试过model.save_weights()加速保存，希望通过修改代码降低模型存储占用。

原实现代码：

class Module_Embedding(tf.keras.Model):

    def __init__(self, num_class_list: List[int], dim_embeddings: int = 4):
        super().__init__()
        self.num_class_list = num_class_list
        self.dim_embeddings = dim_embeddings
        self.dim_ft_categorical = len(num_class_list)
        self.embeddings = [layers.Embedding(nc, self.dim_embeddings, embeddings_regularizer=tf.keras.regularizers.l2(0.2)) for nc in self.num_class_list]

    def call(self, inputs):
        ft_categorical_embed_list = [self.embeddings[i](inputs[:,i]) for i in range(self.dim_ft_categorical)]
        ft_categorical_embed =  tf.concat(ft_categorical_embed_list, axis=-1)
        return ft_categorical_embed


num_features=3000
num_class_list = [10 for _ in range(num_features)]
inputs = tf.random.uniform((n, m), maxval=min(num_class_list))
model = Module_Embedding(num_class_list=num_class_list, dim_embeddings=3)
y = model(inputs)
model.summary()
model.save('./2_model/temp/model_check_0730')

问题原因 #

模型保存体积过大并非参数本身导致，而是大量独立Embedding层带来的计算图元数据冗余。TensorFlow保存SavedModel时，每个layers.Embedding都会生成独立的计算图节点、配置信息和元数据，3000个层会产生巨量重复的结构描述，这才是saved_model.pb文件膨胀的核心原因。

解决方案：合并Embedding层 #

将所有特征的Embedding合并为单个大Embedding层，通过类别ID偏移实现不同特征的独立编码，既保持原功能，又大幅减少计算图元数据。

实现思路 #

1. 计算每个特征的类别ID偏移量：给每个特征的类别ID加上前面所有特征的类别总数，让所有特征的ID变为全局唯一值（如第一个特征0-9，第二个10-19，第三个20-29...）。
2. 使用单个Embedding层，输入维度设为所有特征的类别总数之和。
3. 在call方法中对输入的每个特征ID添加对应偏移量，统一传入大Embedding层，最后将结果展平拼接。

修改后代码 #

from typing import List
import tensorflow as tf
from tensorflow.keras import layers

class Module_Embedding(tf.keras.Model):
    def __init__(self, num_class_list: List[int], dim_embeddings: int = 4):
        super().__init__()
        self.num_class_list = num_class_list
        self.dim_embeddings = dim_embeddings
        self.dim_ft_categorical = len(num_class_list)
        
        # 计算每个特征的类别ID偏移量
        self.offsets = tf.cumsum([0] + num_class_list[:-1], dtype=tf.int32)
        # 合并为单个大Embedding层
        total_classes = sum(num_class_list)
        self.embedding = layers.Embedding(
            input_dim=total_classes,
            output_dim=dim_embeddings,
            embeddings_regularizer=tf.keras.regularizers.l2(0.2)
        )

    def call(self, inputs):
        # 为每个特征的ID添加对应偏移量
        inputs_with_offset = inputs + self.offsets[tf.newaxis, :]
        # 统一通过大Embedding层获取编码
        embeddings = self.embedding(inputs_with_offset)
        # 展平为(batch_size, num_features*dim_embeddings)，与原输出格式一致
        ft_categorical_embed = tf.reshape(embeddings, (-1, self.dim_ft_categorical * self.dim_embeddings))
        return ft_categorical_embed


# 测试代码
num_features = 3000
num_class_list = [10 for _ in range(num_features)]
n, m = 32, num_features  # 补充原代码中缺失的n、m定义
inputs = tf.random.uniform((n, m), maxval=min(num_class_list), dtype=tf.int32)
model = Module_Embedding(num_class_list=num_class_list, dim_embeddings=3)
y = model(inputs)
model.summary()
model.save('./2_model/temp/model_check_0730_optimized')

效果说明 #

修改后模型的参数总数与原实现完全一致，但仅包含1个Embedding层，计算图的元数据量大幅减少，保存后的saved_model.pb文件体积会接近理论参数存储大小（约0.34M），解决了原问题。

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