Method cgc_model

modelzoo/ple/train.py:195–248

(self, inputs, level_name, is_last=False)

    # single Extraction Layer
    def cgc_model(self, inputs, level_name, is_last=False):
        specific_expert_outputs = []
        # build task-specific expert layer
        for i in range(self._num_tasks):
            for j in range(self._specific_expert_num):
                expert_network = self._dnn(inputs[i], dnn_hidden_units=self._expert_dnn_hidden_units, layer_name=level_name + 'task_' + self._towers[i][0] + '_expert_specific_' + str(j))
                specific_expert_outputs.append(expert_network)

        # build task-shared expert layer
        shared_expert_outputs = []
        for k in range(self._shared_expert_num):
            expert_network = self._dnn(inputs[-1], dnn_hidden_units=self._expert_dnn_hidden_units, layer_name=level_name + 'expert_shared_' + str(k))
            shared_expert_outputs.append(expert_network)

        # task_specific gate (count = num_tasks)
        cgc_outs = []
        for i in range(self._num_tasks):
            # concat task-specific expert and task-shared expert
            cur_expert_num = self._specific_expert_num + self._shared_expert_num
            # task_specific + task_shared
            cur_experts = specific_expert_outputs[
                          i * self._specific_expert_num:(i + 1) * self._specific_expert_num] + shared_expert_outputs

            expert_concat = tf.keras.layers.Lambda(lambda x: tf.stack(x, axis=1))(cur_experts)

            # build gate layers
            gate_input = self._dnn(inputs[i], dnn_hidden_units=self._gate_dnn_hidden_units, layer_name=level_name + 'gate_specific_' + self._towers[i][0])
            gate_out = tf.layers.dense(gate_input, units=cur_expert_num,
                                             name=level_name + 'gate_softmax_specific_' + self._towers[i][0])
            gate_out = tf.keras.layers.Lambda(lambda x: tf.expand_dims(x, axis=-1))(gate_out)

            gate_mul_expert = tf.keras.layers.Lambda(lambda x: tf.math.reduce_sum(x[0] * x[1], axis=1, keep_dims=False),
                                                     name=level_name + 'gate_mul_expert_specific_' + self._towers[i][0])(
                [expert_concat, gate_out])
            cgc_outs.append(gate_mul_expert)

        # if not last, add a shared gate
        if not is_last:
            cur_expert_num = self._num_tasks * self._specific_expert_num + self._shared_expert_num
            cur_experts = specific_expert_outputs + shared_expert_outputs  # all the expert include task-specific expert and task-shared expert

            expert_concat = tf.keras.layers.Lambda(lambda x: tf.stack(x, axis=1))(cur_experts)
            # gate layers
            gate_input = self._dnn(inputs[-1], dnn_hidden_units=self._gate_dnn_hidden_units, layer_name=level_name + 'gate_shared')
            gate_out = tf.layers.dense(gate_input, units=cur_expert_num, use_bias=False, activation='softmax',
                                             name=level_name + 'gate_softmax_shared')
            gate_out = tf.keras.layers.Lambda(lambda x: tf.expand_dims(x, axis=-1))(gate_out)

            gate_mul_expert = tf.keras.layers.Lambda(lambda x: tf.math.reduce_sum(x[0] * x[1], axis=1, keep_dims=False),
                                                     name=level_name + 'gate_mul_expert_shared')(
                [expert_concat, gate_out])

            cgc_outs.append(gate_mul_expert)
        return cgc_outs

    # create model
    def _create_model(self):