Method _create_model

modelzoo/dlrm/train.py:135–234

(self)

    # create model
    def _create_model(self):
        # input dense feature and embedding of sparse features
        with tf.variable_scope('input_layer', reuse=tf.AUTO_REUSE):
            with tf.variable_scope('dense_input_layer',
                                   partitioner=self._input_layer_partitioner,
                                   reuse=tf.AUTO_REUSE):
                dense_inputs = tf.feature_column.input_layer(
                    self._feature, self._dense_column)
            with tf.variable_scope('sparse_input_layer', reuse=tf.AUTO_REUSE):
                column_tensors = {}
                if self._adaptive_emb and not self.tf:
                    '''Adaptive Embedding Feature Part 1 of 2'''
                    adaptive_mask_tensors = {}
                    for col in CATEGORICAL_COLUMNS:
                        adaptive_mask_tensors[col] = tf.ones([args.batch_size],
                                                             tf.int32)
                    sparse_inputs = tf.feature_column.input_layer(
                        features=self._feature,
                        feature_columns=self._sparse_column,
                        cols_to_output_tensors=column_tensors,
                        adaptive_mask_tensors=adaptive_mask_tensors)
                else:
                    sparse_inputs = tf.feature_column.input_layer(
                        features=self._feature,
                        feature_columns=self._sparse_column,
                        cols_to_output_tensors=column_tensors)

        # MLP behind dense inputs
        mlp_bot_scope = tf.variable_scope(
            'mlp_bot_layer',
            partitioner=self._dense_layer_partitioner,
            reuse=tf.AUTO_REUSE)
        with mlp_bot_scope.keep_weights(dtype=tf.float32) if self.bf16 \
                else mlp_bot_scope:
            if self.bf16:
                dense_inputs = tf.cast(dense_inputs, dtype=tf.bfloat16)

            for layer_id, num_hidden_units in enumerate(self._mlp_bot):
                with tf.variable_scope(
                        'mlp_bot_hiddenlayer_%d' % layer_id,
                        reuse=tf.AUTO_REUSE) as mlp_bot_hidden_layer_scope:
                    dense_inputs = tf.layers.dense(
                        dense_inputs,
                        units=num_hidden_units,
                        activation=tf.nn.relu,
                        name=mlp_bot_hidden_layer_scope)
                    dense_inputs = tf.layers.batch_normalization(
                        dense_inputs,
                        training=self.is_training,
                        trainable=True)
                    self._add_layer_summary(dense_inputs,
                                            mlp_bot_hidden_layer_scope.name)
            if self.bf16:
                dense_inputs = tf.cast(dense_inputs, dtype=tf.float32)

        # interaction_op
        if self.interaction_op == 'dot':
            # dot op