Method _create_model

modelzoo/dssm/train.py:154–225

(self)

    # create model
    def _create_model(self):
        # input embeddings of user & item features
        with tf.variable_scope('input_layer',
                               partitioner=self._input_layer_partitioner,
                               reuse=tf.AUTO_REUSE):
            if self._adaptive_emb and not self.tf:
                '''Adaptive Embedding Feature Part 1 of 2'''
                adaptive_mask_tensors = {}
                for col in INPUT_FEATURES:
                    adaptive_mask_tensors[col] = tf.ones([self._batch_size],
                                                         tf.int32)
                user_emb = tf.feature_column.input_layer(
                    self._feature,
                    self._user_column,
                    adaptive_mask_tensors=adaptive_mask_tensors)
                item_emb = tf.feature_column.input_layer(
                    self._feature,
                    self._item_column,
                    adaptive_mask_tensors=adaptive_mask_tensors)
            else:
                for key in TAG_COLUMN:
                    self._feature[key] = tf.strings.split(
                        self._feature[key], '|')
                user_emb = tf.feature_column.input_layer(
                    self._feature, self._user_column)
                item_emb = tf.feature_column.input_layer(
                    self._feature, self._item_column)

        if self.bf16:
            user_emb = tf.cast(user_emb, dtype=tf.bfloat16)
            item_emb = tf.cast(item_emb, dtype=tf.bfloat16)

        # user dnn network
        user_scope = tf.variable_scope('user_dnn_layer', \
            partitioner=self._dense_layer_partitioner,
            reuse=tf.AUTO_REUSE)
        with user_scope.keep_weights(dtype=tf.float32) if self.bf16 \
            else user_scope:
            user_emb = self._dnn_tower(user_emb, 'user_dnn')

        # item dnn network
        item_scope = tf.variable_scope('item_dnn_layer', \
            partitioner=self._dense_layer_partitioner,
            reuse=tf.AUTO_REUSE)
        with item_scope.keep_weights(dtype=tf.float32) if self.bf16 \
            else item_scope:
            item_emb = self._dnn_tower(item_emb, 'item_dnn')

        if self.bf16:
            user_emb = tf.cast(user_emb, dtype=tf.float32)
            item_emb = tf.cast(item_emb, dtype=tf.float32)

        # norm
        user_emb = tf.math.l2_normalize(user_emb, axis=1)
        item_emb = tf.math.l2_normalize(item_emb, axis=1)

        user_item_sim = tf.reduce_sum(tf.multiply(user_emb, item_emb),
                                      axis=1,