multi_task_model.py

# -*- coding: utf-8 -*-
import tensorflow as tf
from tensorflow.contrib.layers.python.layers import batch_norm as batch_norm
from utils.tf_utils import *
from functools import reduce



class MultiTaskClassifier:

    def __init__(self, x_dim, sparse_x_len, emb_size, tower_units=None, tower_activations=None, lr=0.001,
                 batch_norm=False, bn_decay=0.9, alpha=1.8, beta=0.5, multi_hot_dict=None,
                 nn_init='glorot_uniform', logger=None, use_sample_weight=False,
                 fl_gamma=0., fl_alpha=1., base_model='fm', n_towers=1, tower_weights='1',
                 expert_units=None, expert_activations=None, n_experts=None, infer_weights=None,
                 export_original_outputs=False, id_col_configs=None, id_seq_col_configs=None,
                 vector_col_dim=None, label_names=None, num_levels=None, keep_pro=None):
        """

        :param x_dim:
        :param sparse_x_len:
        :param emb_size: fm embedding size
        :param tower_units: tower nn units
        :param tower_activations: deepfm nn activations
        :param lr: learning rate
        :param batch_norm: 是否进行batch normalization
        :param bn_decay: batch normalization decay参数
        :param alpha: 正则化参数, 总权重
        :param beta: 正则化参数, L1/L2占比
        :param multi_hot_dict: 标明哪些field是multi-hot的 {field_index: field_size}
        :param nn_init: nn参数初始化方法
        :param logger:
        :param use_sample_weight: 是否使用sample weight
        :param fl_gamma: focal loss的gamma值
        :param fl_alpha: focal loss的alpha值
        :param base_model: 底层模型种类, 如fm, moe, mmoe (moe, mmoe时专家相关参数才生效)
        :param n_towers
        :param tower_weights
        :param expert_units: 专家网络神经元个数
        :param expert_activations: 专家网络激活函数
        :param n_experts: 专家网络个数
        :param infer_weights: 推断时的tower weight
        :param export_original_outputs: 导出模型时是否以所有目标原始值的形式导出
        :param id_col_configs:
        :param vector_col_dim:
        """
        # logger
        self.logger = logger

        # 模型参数设置
        self.x_dim = x_dim
        self.sparse_x_len = sparse_x_len
        self.emb_size = emb_size
        self.tower_units = tower_units
        self.tower_activations = tower_activations
        if self.tower_units is not None and self.tower_activations is not None:
            assert len(self.tower_units) == len(self.tower_activations), 'layer units and activations length not match: %s vs %s' % (len(self.tower_units), len(self.tower_activations))
        self.lr = lr
        self.batch_norm = batch_norm
        self.batch_norm_decay = bn_decay
        self.alpha = alpha
        self.beta = beta
        self.multi_hot_dict = multi_hot_dict
        self.nn_init = nn_init
        self.use_sample_weight = use_sample_weight
        self.mode = None
        self.expert_units = expert_units
        self.expert_activations = expert_activations
        self.n_experts = n_experts
        # multi-hot在输入x中对应的下标
        i_multi_hot = []
        for key in multi_hot_dict:
            temp = list(range(key, key + multi_hot_dict[key]))
            i_multi_hot.extend(temp)

        # one-hot在输入x中对应的下标
        self.i_one_hot = []
        for i in range(sparse_x_len):
            if i not in i_multi_hot:
                self.i_one_hot.append(i)

        self.indices, self.values = [], []
        self.fl_gamma = fl_gamma
        self.fl_alpha = fl_alpha
        self.base_model = base_model
        self.n_towers = n_towers
        self.tower_weights = [float(x) for x in tower_weights.split(',')]
        self.infer_weights = self.tower_weights if infer_weights is None else [float(x) for x in infer_weights.split(',')]
        self.export_original_outputs = export_original_outputs
        self.id_col_configs = id_col_configs
        self.id_seq_col_configs = id_seq_col_configs
        self.keep_pro = keep_pro ## drop out

        self.total_id_embedding_dim = 0
        self.total_id_seq_embedding_dim = 0

        if self.id_col_configs is not None and len(self.id_col_configs) > 0:
            self.total_id_embedding_dim = sum([v['dim'] for v in self.id_col_configs.values()])
        elif self.id_seq_col_configs is not None and len(self.id_seq_col_configs) > 0:
            self.total_id_seq_embedding_dim = sum([v['dim'] for v in self.id_seq_col_configs.values()])
        print_info('total id embedding dim: %s' % self.total_id_embedding_dim, self.logger)
        print_info('total seq id embedding dim: %s'%self.total_id_seq_embedding_dim, self.logger)

        self.vector_col_dim = 0 if vector_col_dim is None else vector_col_dim
        print_info('embedding col dim: %s' % self.vector_col_dim)

        # PLE
        self.label_names = label_names
        self.num_levels = num_levels

        # 模型重要tensor声明
        self.x = None
        self.y = None
        self.pred = None
        self.weighted_pred = None
        self.loss = None
        self.global_step = None
        self.train_op = None
        self.ordr1_weights = None
        self.ordr2_emb_mat = None
        self.sample_weight = None
        self.tower_sample_losses = None
        self.reg_loss = None
        self.param_num = None
        self.summary = None
        self.batch_norm_layer = None
        self.id_map = None
        self.id_embedding = None
        self.vector_x = None
        self.id_seq_map = None
        self.id_seq_embedding = None

    def train_bn_layer(self, x, scope_bn):
        return batch_norm(x, decay=self.batch_norm_decay, center=True, scale=True, is_training=True, reuse=None,
                          trainable=True, scope=scope_bn)

    def inf_bn_layer(self, x, scope_bn):
        return batch_norm(x, decay=self.batch_norm_decay, center=True, scale=True, is_training=False, reuse=None,
                          trainable=False, scope=scope_bn)

    def base_network(self, x, is_train):
        """
        构建多目标模型的底层共享网络
        :param x:
        :param is_train:
        :return: base output tensor
        """
        # 定义batch_norm layer
        if self.batch_norm:
            if is_train:
                self.batch_norm_layer = self.train_bn_layer
            else:
                self.batch_norm_layer = self.inf_bn_layer

        if self.base_model == 'fm':
            # fm前向过程
            # factorization machine只支持二分类
            with tf.variable_scope('var/fm_first_order', reuse=tf.AUTO_REUSE):
                self.ordr1_weights = tf.get_variable('weight_matrix',
                                                     shape=[self.x_dim, 1],
                                                     dtype=tf.float32,
                                                     initializer=tf.initializers.random_normal(stddev=0.01))
                # # 使用最大的index来作为空值
                # ordr1_weights = tf.concat([self.ordr1_weights, tf.zeros(shape=[1, 1])], axis=0)  # [x_dim + 1, 1]
                # ^^其实没必要, 当lookup的id大于lookup param维度时, tf会直接返回0

                # 查找每个离散值index对应维度的order1权重
                first_ordr = tf.nn.embedding_lookup(self.ordr1_weights, x)  # [None, sparse_x_len, 1]
                first_ordr = tf.reduce_sum(first_ordr, 2)   # [None, sparse_x_len], deepFM中用于拼接的一阶向量

                intersect = tf.get_variable('intersect', shape=[1], dtype=tf.float32,
                                            initializer=tf.zeros_initializer())
                # intersect重复batch_size次后进行拼接
                intersect = tf.tile(intersect[tf.newaxis, :], [tf.shape(first_ordr)[0], 1])     # [None, 1]
                first_ordr = tf.concat([first_ordr, intersect], axis=1)     # [None, sparse_x_len + 1] 截距项也加入一阶向量

            with tf.variable_scope('var/fm_second_order', reuse=tf.AUTO_REUSE):
                self.ordr2_emb_mat = tf.get_variable('embedding_matrix',
                                                     shape=[self.x_dim, self.emb_size],
                                                     dtype=tf.float32,
                                                     initializer=tf.initializers.random_normal(stddev=0.01))
                # 使用最大的index来作为空值
                # ordr2_emb_mat = tf.concat([self.ordr2_emb_mat, tf.zeros(shape=[1, self.emb_size])],
                #                           axis=0)  # [x_dim + 1, emb_size]
                dis_embs = tf.nn.embedding_lookup(self.ordr2_emb_mat, x)  # [None, sparse_x_len, emb_size]
                # sum -> sqr part
                sum_dis_embs = tf.reduce_sum(dis_embs, 1)  # [None, emb_size]
                sum_sqr_dis_embs = tf.square(sum_dis_embs)  # [None, emb_size]
                # sqr -> sum part
                sqr_dis_embs = tf.square(dis_embs)  # [None, sparse_x_len, emb_size]
                sqr_sum_dis_embs = tf.reduce_sum(sqr_dis_embs, 1)  # [None, emb_size]
                # second order
                second_ordr = 0.5 * (sum_sqr_dis_embs - sqr_sum_dis_embs)  # [None, emb_size]

            out_tensors = [first_ordr, second_ordr]      # shape = [None, sparse_x_len + 1 + emb_size]
            print(tf.concat(out_tensors, 1).shape)
            if self.id_embedding is not None:
                # 若存在id embedding, 则拼接到输出tensor上
                out_tensors.append(self.id_embedding)    # shape[1] + total_id_embedding_size
                print(tf.concat(out_tensors, 1).shape)

            if self.id_seq_embedding is not None:
                out_tensors.append(self.id_seq_embedding)
                print(tf.concat(out_tensors, 1).shape)

            if self.vector_x is not None:
                # 若存在embedding特征, 则拼接到输出tensor上
                out_tensors.append(self.vector_x)     # shape[1] + vector_x_size
                print(tf.concat(out_tensors, 1).shape)
            return tf.concat(out_tensors, 1)
        elif self.base_model in ('moe', 'mmoe'):
            # 对于moe, 返回tower间共用的一个专家输出tensor
            # 对于mmoe, 返回每个tower自定义的专家输出tensor list
            with tf.variable_scope('var/base_embedding', reuse=tf.AUTO_REUSE):
                self.ordr2_emb_mat = tf.get_variable('embedding_matrix',
                                                     shape=[self.x_dim, self.emb_size],
                                                     dtype=tf.float32,
                                                     initializer=tf.initializers.random_normal(stddev=0.01))
                dis_embs = tf.nn.embedding_lookup(self.ordr2_emb_mat, x)  # [None, sparse_x_len, emb_size]
                # 通过embedding将输入维度从x_dim降至sparse_x_len * emb_size
                dense_x = [tf.reshape(dis_embs, shape=(-1, self.sparse_x_len * self.emb_size))]
                print(tf.concat(dense_x, 1).shape)
                if self.id_embedding is not None:
                    # 若存在id embedding, 则拼接到第一层神经网络的输入上
                    dense_x.append(self.id_embedding)       # shape[1] + total_id_embedding_size
                    print(tf.concat(dense_x, 1).shape)
                if self.vector_x is not None:
                    # 若存在embedding特征, 则拼接到第一层神经网络的输入上
                    dense_x.append(self.vector_x)        # shape[1] + vector_x_size
                    print(tf.concat(dense_x, 1).shape)
                dense_x = tf.concat(dense_x, 1)

            with tf.variable_scope('var/expert_networks', reuse=tf.AUTO_REUSE):
                # 构建专家网络
                experts = []
                for i in range(self.n_experts):
                    # 专家网络输出向量shape=(None, expert_units[-1])
                    out = get_mlp_network(
                        dense_x,
                        self.expert_units,
                        self.expert_activations,
                        'expert_%s' % i,
                        bn_layer=self.batch_norm_layer,
                        disable_last_bn=False
                    )
                    out = out[:, tf.newaxis, :]                             # shape=(None, 1, expert_units[-1])
                    experts.append(out)
                experts = tf.concat(experts, 1, name='expert_origin_outputs')      # shape=(None, n_experts, expert_units[-1])

                # 专家控制门
                if self.base_model == 'moe':
                    # 对于moe, 所有tower共享一个gate
                    gate_kernel = tf.get_variable('expert_gate',
                                                  shape=[self.sparse_x_len * self.emb_size + self.total_id_embedding_dim + self.vector_col_dim, self.n_experts],
                                                  dtype=tf.float32,
                                                  initializer=tf.initializers.random_normal(stddev=0.01))
                    gate = tf.matmul(dense_x, gate_kernel)  # shape=(None, n_experts)
                    gate = tf.nn.softmax(gate)
                    gate = tf.tile(gate[:, :, tf.newaxis], [1, 1, self.expert_units[-1]])  # shape=(None, n_experts, expert_units[-1])

                    # 专家网络输出根据控制门进行加权得到最终输出
                    out = experts * gate                            # shape=(None, n_experts, expert_units[-1])
                    experts_final = tf.reduce_sum(out, axis=1)      # shape=(None, experts_units[-1])
                else:
                    # 对于mmoe, 每个tower都有其单独的gate和公共的专家网络进行加权, 循环多次为每个tower计算控制门
                    experts_final = []
                    for i in range(self.n_towers):
                        gate_kernel = tf.get_variable('expert_gate_%s' % i,
                                                      shape=[self.sparse_x_len * self.emb_size + self.total_id_embedding_dim + self.vector_col_dim, self.n_experts],
                                                      dtype=tf.float32,
                                                      initializer=tf.initializers.random_normal(stddev=0.01))
                        gate = tf.matmul(dense_x, gate_kernel)                      # shape=(None, n_experts)
                        gate = tf.nn.softmax(gate)
                        gate = tf.tile(gate[:, :, tf.newaxis], [1, 1, self.expert_units[-1]])   # shape=(None, n_experts, expert_units[-1])

                        # 专家网络输出根据控制门进行加权得到最终输出
                        out = experts * gate                              # shape=(None, n_experts, expert_units[-1])
                        out = tf.reduce_sum(out, axis=1)                  # shape=(None, experts_units[-1])
                        experts_final.append(out)

                return experts_final
        elif self.base_model == 'ple':
            with tf.variable_scope('var/base_embedding', reuse=tf.AUTO_REUSE):
                self.ordr2_emb_mat = tf.get_variable('embedding_matrix',
                                                     shape=[self.x_dim, self.emb_size],
                                                     dtype=tf.float32,
                                                     initializer=tf.initializers.random_normal(stddev=0.01))
                dis_embs = tf.nn.embedding_lookup(self.ordr2_emb_mat, x)  # [None, sparse_x_len, emb_size]
                # 通过embedding将输入维度从x_dim降至sparse_x_len * emb_size
                dense_x = [tf.reshape(dis_embs, shape=(-1, self.sparse_x_len * self.emb_size))]
                print(tf.concat(dense_x, 1).shape)
                if self.id_embedding is not None:
                    # 若存在id embedding, 则拼接到第一层神经网络的输入上
                    dense_x.append(self.id_embedding)  # shape[1] + total_id_embedding_size
                    print(tf.concat(dense_x, 1).shape)

                if self.id_seq_embedding is not None:
                    dense_x.append(self.id_seq_embedding)
                    print(tf.concat(dense_x, 1).shape)

                if self.vector_x is not None:
                    # 若存在embedding特征, 则拼接到第一层神经网络的输入上
                    dense_x.append(self.vector_x)  # shape[1] + vector_x_size
                    print(tf.concat(dense_x, 1).shape)
                dense_x = tf.concat(dense_x, 1)

            with tf.variable_scope('var/extract_network', reuse=tf.AUTO_REUSE):
                outputs = dense_x
                for level in range(self.num_levels):
                    # 如果不是第一层，那么输入是多个上层的输出expert组成的列表
                    # 此时，需要进行fusion：一般是拼接、相乘、相加几种融合方式
                    # 这里使用相加拼接相乘
                    if isinstance(outputs, list):
                        outputs = tf.concat([reduce(lambda x, y: x + y, outputs),
                                             reduce(lambda x, y: x * y, outputs)],
                                            axis=-1)

                    with tf.variable_scope('Mixture-of-Experts', reuse=tf.AUTO_REUSE):
                        mixture_experts = []
                        for name in self.label_names + ["expert_shared"]:
                            for i in range(self.n_experts):
                                # expert一般是一层全连接层
                                out = get_mlp_network(
                                    input_tensor=outputs,
                                    layer_units=[self.expert_units[level]],
                                    layer_activations=[self.expert_activations[level]],
                                    name="{}_expert_{}_level_{}".format(name, i, level),
                                    bn_layer=self.batch_norm_layer,
                                    disable_last_bn=False
                                )

                                # out = out[:, tf.newaxis, :]  # shape=(None, 1, expert_units[-1])
                                mixture_experts.append(out)
                    # 如果是最后一层，那么gate的数量应该是task的数量
                    # 其他层的话，gate的数量一般等于experts的数量
                    if level == self.num_levels - 1:
                        num_gates = len(self.label_names)
                    else:
                        num_gates = self.n_experts

                    # 生成不同专家任务的门
                    with tf.variable_scope("multi-gate", reuse=tf.AUTO_REUSE):
                        multi_gate = []
                        for i in range(num_gates):
                            gate = tf.layers.dense(dense_x, units=self.n_experts * (len(self.label_names) + 1),
                                                   kernel_initializer=tf.initializers.random_uniform(-0.01, 0.01),
                                                   bias_initializer=tf.initializers.zeros(),
                                                   name="gate_{}_level{}".format(i, level)
                                                   )
                            gate = tf.nn.softmax(gate)
                            multi_gate.append(gate)

                    with tf.variable_scope("combine_gate_expert", reuse=tf.AUTO_REUSE):
                        ple_layers = []
                        for i in range(num_gates):
                            ple_layers.append(self._combine_expert_gate(mixture_experts, multi_gate[i]))
                    outputs = ple_layers
                return outputs
        else:
            raise TypeError('不支持的base_model类型: %s' % self.base_model)

    def _combine_expert_gate(self, mixture_experts, gate):
        mixture_experts = tf.concat([tf.expand_dims(dnn, axis=1) for dnn in mixture_experts], axis=1)
        gate = tf.expand_dims(gate, axis=-1)
        return tf.reduce_sum(mixture_experts * gate, axis=-1)

    def tower_network(self, base_output, is_train):
        """
        构建各个目标的专用网络
        :param base_output: tensor or list<tensor>
        :param is_train:
        :return:
        """
        with tf.variable_scope('var/tower_networks', reuse=tf.AUTO_REUSE):
            # 定义batch_norm layer
            if self.batch_norm:
                if is_train:
                    self.batch_norm_layer = self.train_bn_layer
                else:
                    self.batch_norm_layer = self.inf_bn_layer

            towers = []
            for i in range(self.n_towers):
                if self.tower_units is None or self.tower_activations is None:
                    # 当不设定tower网络参数时, 直接将输入进来的tensor进行reduce_sum后输出(兼容fm)
                    print('no tower network')
                    print(base_output.shape)
                    print(tf.reduce_sum(base_output, axis=1, keepdims=True).shape)
                    towers.append(tf.reduce_sum(base_output, axis=1, keepdims=True))
                else:
                    # 若base_output为单一tensor, 则每个tower都直接使用这一tensor; 若为list, 则每个tower选取使用对应的tensor
                    tower_input = base_output if type(base_output) is not list else base_output[i]
                    tower_out = get_mlp_network(
                        tower_input,
                        self.tower_units,
                        self.tower_activations,
                        'tower_%s' % i,
                        bn_layer=self.batch_norm_layer,
                        disable_last_bn=True)
                    towers.append(tower_out)
            return tf.concat(towers, 1, name='tower_outputs')

    def forward(self, is_train):
        """

        :param is_train:
        :return:
        """
        # 对id类型的输入进行embedding
        if self.id_map is not None and len(self.id_map) > 0:
            with tf.variable_scope('var/id_embedding', reuse=tf.AUTO_REUSE):
                columns = []
                for id_name in self.id_map.keys():
                    # id --> new hash id
                    id_config = self.id_col_configs[id_name]
                    hash_column = tf.feature_column.categorical_column_with_hash_bucket(id_name,
                                                                                        id_config['size'],
                                                                                        dtype=tf.int32)
                    # new hash id --> embedding
                    embedding_column = tf.feature_column.embedding_column(hash_column,
                                                                          id_config['dim'],
                                                                          trainable=True)
                    columns.append(embedding_column)
                self.id_embedding = tf.feature_column.input_layer(self.id_map, columns)    # shape=(None, sum dim of all ids)

        if self.id_seq_map is not None and len(self.id_seq_map) > 0:
            with tf.variable_scope('var/id_seq_embedding', reuse=tf.AUTO_REUSE):
                columns = []
                for id_name in self.id_seq_map.keys():
                    id_seq_config = self.id_seq_col_configs[id_name]
                    seq_hash_column = tf.contrib.feature_column.sequence_categorical_column_with_hash_bucket(id_name,
                                                                                                             id_seq_config[
                                                                                                                 'size'],
                                                                                                             dtype=tf.int32)
                    seq_embedding_column = [tf.feature_column.embedding_column(hash_column,
                                                                               dimension=id_seq_config['dim'],
                                                                               trainable=True) for hash_column in
                                            seq_hash_column]
                    columns.extend(seq_embedding_column)
                self.id_seq_embedding = tf.feature_column.input_layer(self.id_seq_map, columns)
                print_info('id seq embedding %s' % (self.id_seq_embedding.shape), self.logger)

        base_output = self.base_network(self.x, is_train)
        return self.tower_network(base_output, is_train)

    def cal_loss(self):
        # shape=[None, n_towers]
        # use focal loss
        tower_sample_losses = focal_loss(
            self.y,
            self.pred,
            gamma=self.fl_gamma,
            alpha=self.fl_alpha
        ) * tf.constant(self.tower_weights)                             # shape=[None, n_towers]
        if self.sample_weight is not None:
            tower_sample_losses *= self.sample_weight                   # shape=[None, n_towers]

        sample_loss = tf.reduce_sum(tower_sample_losses, axis=1)        # shape=[None,]
        sample_loss = tf.reduce_mean(sample_loss)                       # shape=[] (scalar)

        # 所有参数组成的向量，用以添加正则项
        var_list = tf.trainable_variables(scope='train/var')
        _var_vector = tf.concat([tf.reshape(v, shape=(-1,)) for v in var_list], axis=0)
        if self.mode == tf.estimator.ModeKeys.TRAIN:
            for v in var_list:
                self.print('reg var: %s' % v)
                self.print('*' * 50)
            self.print('参数量: %s' % _var_vector.shape)
        self.param_num = _var_vector.shape
        reg_loss = self.alpha * (self.beta * tf.reduce_sum(tf.abs(_var_vector))
                                 + (1 - self.beta) * tf.reduce_sum(tf.square(_var_vector)))
        loss = sample_loss + reg_loss
        return loss, tower_sample_losses, reg_loss

    def build_model(self, x=None, vector_x=None, id_map=None, id_seq_map=None, y=None, w=None, mode=None, hooks=None):
        """

        :param x: tensor类型, 稀疏表达的multi-hot向量, 包含所有向量值为1的indices
        :param vector_x: tensor类型, embedding类型的特征
        :param id_map: dict类型, e.g. {'uid': tensor}, 需要进行embedding的id特征
        :param y: tensor类型, label
        :param w: tensor类型, sample weight
        :param mode: tf.estimator.ModeKeys的某个枚举值, 用于指定模型处于训练、评估或是推断模式
        :param hooks: TensorFlow hook
        :return:
        """
        self.mode = mode
        self.global_step = tf.train.get_or_create_global_step()

        with tf.variable_scope('input', reuse=tf.AUTO_REUSE):
            self.x = tf.reshape(x, shape=(-1, self.sparse_x_len))   # 任何情况下x都不能为空
            self.id_map = id_map
            self.vector_x = vector_x
            self.id_seq_map = id_seq_map

            if y is not None:
                # 对于推理过程, y可以为空
                self.y = tf.cast(tf.reshape(y, (-1, self.n_towers)), dtype=tf.float32)

        with tf.variable_scope('train', reuse=tf.AUTO_REUSE):
            if mode == tf.estimator.ModeKeys.TRAIN:
                self.pred = self.forward(is_train=True)     # shape=[None, n_towers]
            else:
                self.pred = self.forward(is_train=False)    # shape=[None, n_towers]

            self.pred = tf.reshape(self.pred, (-1, self.n_towers))
            self.weighted_pred = tf.reduce_sum(self.pred * tf.constant(self.infer_weights), axis=1, keepdims=True)     # shape=[None, 1]

            if mode == tf.estimator.ModeKeys.PREDICT:
                # 对于推理过程, 不应计算loss, 否则会强制要求输入y
                pass
            else:
                # 非推理过程, 需要计算loss
                self.sample_weight = w if self.use_sample_weight else None
                self.loss, self.tower_sample_losses, self.reg_loss = self.cal_loss()
                self.summary = tf.summary.scalar('loss', self.loss)

            update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
            with tf.control_dependencies(update_ops):
                if mode == tf.estimator.ModeKeys.TRAIN:
                    # 训练过程
                    self.train_op = tf.train.AdamOptimizer(self.lr) \
                        .minimize(self.loss, global_step=self.global_step)
                    spec = tf.estimator.EstimatorSpec(
                        mode=mode,
                        loss=self.loss,
                        train_op=self.train_op,
                        training_hooks=hooks
                    )
                    tf.summary.scalar('loss', self.loss)
                elif mode == tf.estimator.ModeKeys.PREDICT:
                    # 推理过程
                    if self.export_original_outputs:
                        out_dict = {'prediction': self.pred}                    # shape=[None, n_towers]
                    else:
                        out_dict = {'prediction': self.weighted_pred}           # shape=[None, 1]
                    export_outputs = {
                        tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
                            tf.estimator.export.PredictOutput(out_dict)}
                    spec = tf.estimator.EstimatorSpec(mode=mode,
                                                      predictions=out_dict,
                                                      export_outputs=export_outputs)
                elif mode == tf.estimator.ModeKeys.EVAL:
                    # 评估过程
                    eval_metric_ops = {
                    }

                    for i in range(self.n_towers):
                        eval_metric_ops['tower_%s_loss' % i] = tf.metrics.mean(self.tower_sample_losses[:, i])
                        eval_metric_ops['tower_%s_auc' % i] = tf.metrics.auc(self.y[:, i], self.pred[:, i])

                    spec = tf.estimator.EstimatorSpec(mode=mode,
                                                      loss=self.loss,
                                                      eval_metric_ops=eval_metric_ops)
                else:
                    raise ValueError('invalid mode: %s' % mode)
                return spec

    def print(self, msg):
        if self.logger is not None:
            self.logger.info(msg)
        else:
            print(msg)


if __name__ == '__main__':
    # 可直接运行, 测试类的逻辑正确性
    # mock data
    batch_size, x_dim, sparse_x_len, emb_size = 3, 100, 10, 5
    n_towers = 3
    target_weight = tf.constant([0.5, 0.2, 0.3])  # shape=(3,)

    tower_units = [32, 32, 1]
    tower_activations = ['relu', 'relu', 'sigmoid']
    # tower_units = None
    # tower_activations = None

    data_x = np.random.random_integers(0, x_dim - 1, size=(batch_size, sparse_x_len))
    data_y = np.random.random_integers(0, 1, size=(batch_size, n_towers))
    data_w = np.random.standard_normal(size=(batch_size, n_towers))
    data_id = {
        'pgc_id': [
            [125112364],
            [124163423],
            [198274982]
        ],
        'uid': [
            [156],
            [563],
            [789]
        ]
    }
    vector_col_dim = 5
    data_vector_x = np.random.standard_normal(size=(batch_size, vector_col_dim))

    # place holders
    placeholder_x = tf.placeholder(dtype=tf.int64, shape=(None, sparse_x_len))
    placeholder_y = tf.placeholder(dtype=tf.int64, shape=(None, n_towers))
    placeholder_w = tf.placeholder(dtype=tf.float32, shape=(None, n_towers))
    placeholder_id = {
        'pgc_id': tf.placeholder(dtype=tf.int64, shape=(None, 1)),
        'uid': tf.placeholder(dtype=tf.int64, shape=(None, 1))
    }
    id_col_configs = {
        'pgc_id': {
            'dim': 10,
            'size': 5
        },
        'uid': {
            'dim': 20,
            'size': 2
        }
    }
    # placeholder_id = None
    # id_col_configs = None
    placeholder_vector_x = tf.placeholder(dtype=tf.float32, shape=[None, 5])

    # model init
    model = MultiTaskClassifier(x_dim, sparse_x_len, emb_size,
                                tower_units=tower_units, tower_activations=tower_activations, lr=0.01,
                                batch_norm=True, use_sample_weight=True, base_model='fm', n_towers=n_towers,
                                tower_weights='0.5,0.2,0.3', multi_hot_dict={},
                                expert_units=[16, 16], expert_activations=['relu', 'relu'], n_experts=5,
                                id_col_configs=id_col_configs, vector_col_dim=vector_col_dim)
    model.build_model(x=placeholder_x, id_map=placeholder_id, vector_x=placeholder_vector_x, y=placeholder_y, w=placeholder_w, mode=tf.estimator.ModeKeys.TRAIN)
    with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())
        sess.run(tf.local_variables_initializer())
        for _ in range(100):
            _, loss = sess.run([model.train_op, model.loss],
                               feed_dict={model.x: data_x, model.vector_x: data_vector_x,
                                          model.y: data_y, model.sample_weight: data_w,
                                          model.id_map['pgc_id']: data_id['pgc_id'],
                                          model.id_map['uid']: data_id['uid']
                                          })
            print(loss)
        print(sess.run(model.pred, feed_dict={model.x: data_x, model.vector_x: data_vector_x,
                                              model.id_map['pgc_id']: data_id['pgc_id'],
                                              model.id_map['uid']: data_id['uid']
                                              }))