MMOE核心代码块

最新推荐文章于 2025-05-20 21:49:15 发布

原创最新推荐文章于 2025-05-20 21:49:15 发布 · 838 阅读

0 ·

CC 4.0 BY-SA版权

文章标签：

#MMOE

推荐系统同时被 2 个专栏收录

14 篇文章

订阅专栏

算法

11 篇文章

订阅专栏

本文介绍了一个基于TensorFlow实现的MMoE（Multiple Experts with Multiple Objective functions）多任务学习层的具体代码。该层包含多个专家网络与门网络，能够通过门控机制为每个任务分配不同的专家权重，实现多任务间的资源共享。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

class MMoE_Layer(tf.keras.layers.Layer):
    def __init__(self,expert_dim,n_expert,n_task):
        super(MMoE_Layer, self).__init__()
        self.n_task = n_task
        self.expert_layer = [Dense(expert_dim,activation = 'relu') for i in range(n_expert)]
        self.gate_layers = [Dense(n_expert,activation = 'softmax') for i in range(n_task)]
    def call(self,x):
        # 构建多个专家网络
        E_net = [expert(x) for expert in self.expert_layer]
        E_net = Concatenate(axis = 1)([e[:,tf.newaxis,:] for e in E_net]) # 维度 (bs,n_expert,n_dims)
        # 构建多个门网络
        gate_net = [gate(x) for gate in self.gate_layers]     # 维度 n_task个(bs,n_expert)
        # towers计算：对应的门网络乘上所有的专家网络
        towers = []
        for i in range(self.n_task):
            g = tf.expand_dims(gate_net[i],axis = -1)  # 维度(bs,n_expert,1)
            _tower = tf.matmul(E_net, g,transpose_a=True)
            towers.append(Flatten()(_tower))           # 维度(bs,expert_dim)
        return towers

核心代码块