Method _create_model

modelzoo/ple/train.py:251–308

(self)

    # create model
    def _create_model(self):
        TAG_COLUMN = ['tag_category_list', 'tag_brand_list']
        for key in TAG_COLUMN:
            self._feature[key] = tf.strings.split(self._feature[key], '|')

        key_dict = {}
        with self._make_scope('input_layer', self._bf16, self._input_layer_partitioner):
            print('Adaptive emb = ', self._adaptive_emb, 'TF = ', self.tf)
            if self._adaptive_emb and not self.tf:
                '''Adaptive Embedding Feature part 1 of 2'''
                print('Adaptive Embedding Feature part 1 of 2')
                adaptive_mask_tensors = {}
                for col in HASH_INPUTS:
                    adaptive_mask_tensors[col] = tf.ones([args.batch_size],
                                                          tf.int32)
                input_emb = tf.feature_column.input_layer(
                    self._feature,
                    self._feature_column,
                    adaptive_mask_tensors=adaptive_mask_tensors,
                    cols_to_output_tensors=key_dict)
            else:
                input_emb = tf.feature_column.input_layer(
                    self._feature,
                    self._feature_column,
                    cols_to_output_tensors=key_dict)

        with self._make_scope('PLE', self._bf16, self._dense_layer_partitioner):
            if self._bf16:
                input_emb = tf.cast(input_emb, dtype=tf.bfloat16)
            ple_inputs = [input_emb] * (self._num_tasks + 1)
            ple_outputs = []
            for i in range(self._num_layers):
                with tf.variable_scope(f'extraction_network_{i}'):
                    if i == self._num_layers - 1: # the last extraction net
                        ple_outputs = self.cgc_model(inputs=ple_inputs, level_name='level_'+str(i)+'_', is_last=True)
                    else:
                        ple_outputs = self.cgc_model(inputs=ple_inputs, level_name='level_'+str(i)+'_', is_last=False)
                        ple_inputs = ple_outputs

            towers=[]
            for i, tower in enumerate(self._towers):
                tower_name = tower[0]
                hidden_units = tower[2]

                with tf.variable_scope(tower_name, reuse=tf.AUTO_REUSE):
                    tower_output = self._dnn(ple_outputs[i], dnn_hidden_units=hidden_units, layer_name='tower_'+tower_name)
                    final_tower_predict = tf.layers.dense(inputs=tower_output,
                                                        units=1,
                                                        activation=None,
                                                        name=f'{tower_name}_output')
                    self._add_layer_summary(final_tower_predict, f'{tower_name}_output')
                    if self._bf16:
                        final_tower_predict = tf.cast(final_tower_predict, dtype=tf.float32)
                    towers.append(final_tower_predict)
            tower_stack = tf.stack(towers, axis=1)
            self._logits = tf.squeeze(tower_stack, [2])
            self.probability = tf.math.sigmoid(self._logits)
            self.output = tf.round(self.probability)