train.py

# Copyright 2024 The DeepRec Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

import time
import argparse
import tensorflow as tf
import os
import sys
import math
import collections
from tensorflow.python.client import timeline
import json

from tensorflow.python.ops import partitioned_variables
from dynamic_embedding_server.python import cluster_utils
import gazer
import dynamic_embedding_server as des
import tf_fault_tolerance as ft

# Set to INFO for tracking training, default is WARN. ERROR for least messages
tf.logging.set_verbosity(tf.logging.INFO)
print("Using TensorFlow version %s" % (tf.__version__))

# Definition of some constants
CONTINUOUS_COLUMNS = ['I' + str(i) for i in range(1, 14)]  # 1-13 inclusive
CATEGORICAL_COLUMNS = ['C' + str(i) for i in range(1, 27)]  # 1-26 inclusive
LABEL_COLUMN = ['clicked']
TRAIN_DATA_COLUMNS = LABEL_COLUMN + CONTINUOUS_COLUMNS + CATEGORICAL_COLUMNS
FEATURE_COLUMNS = CONTINUOUS_COLUMNS + CATEGORICAL_COLUMNS
HASH_BUCKET_SIZES = {
    'C1': 2500,
    'C2': 2000,
    'C3': 5000000,
    'C4': 1500000,
    'C5': 1000,
    'C6': 100,
    'C7': 20000,
    'C8': 4000,
    'C9': 20,
    'C10': 100000,
    'C11': 10000,
    'C12': 5000000,
    'C13': 40000,
    'C14': 100,
    'C15': 100,
    'C16': 3000000,
    'C17': 50,
    'C18': 10000,
    'C19': 4000,
    'C20': 20,
    'C21': 4000000,
    'C22': 100,
    'C23': 100,
    'C24': 250000,
    'C25': 400,
    'C26': 100000
}


class DeepFM():
    def __init__(self,
                 wide_column=None,
                 fm_column=None,
                 deep_column=None,
                 dnn_hidden_units=[1024, 256, 32],
                 final_hidden_units=[128, 64],
                 optimizer_type='adam',
                 learning_rate=0.001,
                 inputs=None,
                 use_bn=True,
                 bf16=False,
                 stock_tf=None,
                 adaptive_emb=False,
                 input_layer_partitioner=None,
                 dense_layer_partitioner=None):
        if not inputs:
            raise ValueError('Dataset is not defined.')
        self._feature = inputs[0]
        self._label = inputs[1]

        self._wide_column = wide_column
        self._deep_column = deep_column
        self._fm_column = fm_column
        if not wide_column or not fm_column or not deep_column:
            raise ValueError(
                'Wide column, FM column or Deep column is not defined.')

        self.tf = stock_tf
        self.bf16 = False if self.tf else bf16
        self.is_training = True
        self.use_bn = use_bn
        self._adaptive_emb = adaptive_emb

        self._dnn_hidden_units = dnn_hidden_units
        self._final_hidden_units = final_hidden_units
        self._optimizer_type = optimizer_type
        self._learning_rate = learning_rate
        self._input_layer_partitioner = input_layer_partitioner
        self._dense_layer_partitioner = dense_layer_partitioner

        self._create_model()
        with tf.name_scope('head'):
            self._create_loss()
            self._create_optimizer()
            self._create_metrics()

    # used to add summary in tensorboard
    def _add_layer_summary(self, value, tag):
        tf.summary.scalar('%s/fraction_of_zero_values' % tag,
                          tf.nn.zero_fraction(value))
        tf.summary.histogram('%s/activation' % tag, value)

    def _dnn(self, dnn_input, dnn_hidden_units=None, layer_name=''):
        for layer_id, num_hidden_units in enumerate(dnn_hidden_units):
            with tf.variable_scope(layer_name + '_%d' % layer_id,
                                   partitioner=self._dense_layer_partitioner,
                                   reuse=tf.AUTO_REUSE) as dnn_layer_scope:
                dnn_input = tf.layers.dense(
                    dnn_input,
                    units=num_hidden_units,
                    activation=tf.nn.relu,
                    name=dnn_layer_scope)
                if self.use_bn:
                    dnn_input = tf.layers.batch_normalization(
                        dnn_input, training=self.is_training, trainable=True)
                self._add_layer_summary(dnn_input, dnn_layer_scope.name)

        return dnn_input

    def _create_model(self):
        # input features
        with tf.variable_scope('input_layer',
                               partitioner=self._input_layer_partitioner,
                               reuse=tf.AUTO_REUSE):
            fm_cols = {}
            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)
                dnn_input = tf.feature_column.input_layer(
                    features=self._feature,
                    feature_columns=self._deep_column,
                    adaptive_mask_tensors=adaptive_mask_tensors)
                wide_input = tf.feature_column.input_layer(
                self._feature, self._wide_column, cols_to_output_tensors=fm_cols, adaptive_mask_tensors=adaptive_mask_tensors)
            else:
                dnn_input = tf.feature_column.input_layer(self._feature,
                                                          self._deep_column)
                wide_input = tf.feature_column.input_layer(
                self._feature, self._wide_column, cols_to_output_tensors=fm_cols)

            fm_input = tf.stack([fm_cols[cols] for cols in self._fm_column], 1)

        if self.bf16:
            wide_input = tf.cast(wide_input, dtype=tf.bfloat16)
            fm_input = tf.cast(fm_input, dtype=tf.bfloat16)
            dnn_input = tf.cast(dnn_input, dtype=tf.bfloat16)

        # DNN part
        dnn_scope = tf.variable_scope('dnn')
        with dnn_scope.keep_weights(dtype=tf.float32) if self.bf16 \
                else dnn_scope:
            dnn_output = self._dnn(dnn_input, self._dnn_hidden_units,
                                   'dnn_layer')

        # linear / fisrt order part
        with tf.variable_scope('linear', reuse=tf.AUTO_REUSE) as linear:
            linear_output = tf.reduce_sum(wide_input, axis=1, keepdims=True)

        # FM second order part
        with tf.variable_scope('fm', reuse=tf.AUTO_REUSE) as fm:
            sum_square = tf.square(tf.reduce_sum(fm_input, axis=1))
            square_sum = tf.reduce_sum(tf.square(fm_input), axis=1)
            fm_output = 0.5 * tf.subtract(sum_square, square_sum)

        # Final dnn layer
        all_input = tf.concat([dnn_output, linear_output, fm_output], 1)
        final_dnn_scope = tf.variable_scope('final_dnn')
        with final_dnn_scope.keep_weights(dtype=tf.float32) if self.bf16 \
                else final_dnn_scope:
            dnn_logits = self._dnn(all_input, self._final_hidden_units, 'final_dnn')

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

        self._logits = tf.layers.dense(dnn_logits, units=1)
        self.probability = tf.math.sigmoid(self._logits)
        self.output = tf.round(self.probability)

    # compute loss
    def _create_loss(self):
        loss_func = tf.losses.mean_squared_error
        predict = tf.squeeze(self.probability)
        self.loss = tf.math.reduce_mean(loss_func(self._label, predict))
        tf.summary.scalar('loss', self.loss)

    # define optimizer and generate train_op
    def _create_optimizer(self):
        self.global_step = tf.train.get_or_create_global_step()
        if self.tf or self._optimizer_type == 'adam':
            optimizer = tf.train.AdamOptimizer(
                learning_rate=self._learning_rate,
                beta1=0.9,
                beta2=0.999,
                epsilon=1e-8)
        elif self._optimizer_type == 'adagrad':
            optimizer = tf.train.AdagradOptimizer(
                learning_rate=self._learning_rate,
                initial_accumulator_value=1e-8)
        elif self._optimizer_type == 'adamasync':
            optimizer = tf.train.AdamAsyncOptimizer(
                learning_rate=self._learning_rate,
                beta1=0.9,
                beta2=0.999,
                epsilon=1e-8)
        elif self._optimizer_type == 'adagraddecay':
            optimizer = tf.train.AdagradDecayOptimizer(
                learning_rate=self._learning_rate,
                global_step=self.global_step)
        else:
            raise ValueError('Optimizer type error.')

        update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
        self.global_step_add = tf.assign_add(self.global_step, 1)
        with tf.control_dependencies(update_ops):
            self.train_op = optimizer.minimize(
                self.loss, global_step=self.global_step)

    # compute acc & auc
    def _create_metrics(self):
        self.acc, self.acc_op = tf.metrics.accuracy(labels=self._label,
                                                    predictions=self.output)
        self.auc, self.auc_op = tf.metrics.auc(labels=self._label,
                                               predictions=self.probability,
                                               num_thresholds=1000)
        tf.summary.scalar('eval_acc', self.acc)
        tf.summary.scalar('eval_auc', self.auc)


# generate dataset pipline
def build_model_input(filename, batch_size, num_epochs):
    def parse_csv(value):
        tf.logging.info('Parsing {}'.format(filename))
        cont_defaults = [[0.0] for i in range(1, 14)]
        cate_defaults = [[' '] for i in range(1, 27)]
        label_defaults = [[0]]
        column_headers = TRAIN_DATA_COLUMNS
        record_defaults = label_defaults + cont_defaults + cate_defaults
        columns = tf.io.decode_csv(value, record_defaults=record_defaults)
        all_columns = collections.OrderedDict(zip(column_headers, columns))
        labels = all_columns.pop(LABEL_COLUMN[0])
        features = all_columns
        return features, labels

    def parse_parquet(value):
        tf.logging.info('Parsing {}'.format(filename))
        labels = value.pop(LABEL_COLUMN[0])
        features = value
        return features, labels

    '''Work Queue Feature'''
    if args.workqueue and not args.tf:
        from tensorflow.python.ops.work_queue import WorkQueue
        work_queue = WorkQueue([filename], num_epochs=num_epochs)
        # For multiple files：
        # work_queue = WorkQueue([filename, filename1,filename2,filename3])
        files = work_queue.input_dataset()
    else:
        files = filename
    # Extract lines from input files using the Dataset API.
    if args.parquet_dataset and not args.tf:
        from tensorflow.python.data.experimental.ops import parquet_dataset_ops
        dataset = parquet_dataset_ops.ParquetDataset(files, batch_size=batch_size)
        if args.parquet_dataset_shuffle:
            dataset = dataset.shuffle(buffer_size=20000,
                                      seed=args.seed)  # fix seed for reproducing
        if not args.workqueue:
            dataset = dataset.repeat(num_epochs)
        dataset = dataset.map(parse_parquet, num_parallel_calls=28)
    else:
        dataset = tf.data.TextLineDataset(files)
        dataset = dataset.shuffle(buffer_size=20000,
                                  seed=args.seed)  # fix seed for reproducing
        if not args.workqueue:
            dataset = dataset.repeat(num_epochs)
        dataset = dataset.batch(batch_size)
        dataset = dataset.map(parse_csv, num_parallel_calls=28)
    dataset = dataset.prefetch(2)
    iterator = dataset.make_one_shot_iterator()
    features, labels = iterator.get_next()
    return features, labels


def build_feature_columns():
    wide_column = []
    deep_column = []
    fm_column = []
    if args.group_embedding and not args.tf:
        with tf.feature_column.group_embedding_column_scope(name="categorical"):
            for column_name in FEATURE_COLUMNS:
                if column_name in CATEGORICAL_COLUMNS:
                    categorical_column = tf.feature_column.categorical_column_with_hash_bucket(
                        column_name,
                        hash_bucket_size=10000,
                        dtype=tf.string)

                    if not args.tf:
                        '''Feature Elimination of EmbeddingVariable Feature'''
                        if args.ev_elimination == 'gstep':
                            # Feature elimination based on global steps
                            evict_opt = tf.GlobalStepEvict(steps_to_live=4000)
                        elif args.ev_elimination == 'l2':
                            # Feature elimination based on l2 weight
                            evict_opt = tf.L2WeightEvict(l2_weight_threshold=1.0)
                        else:
                            evict_opt = None
                        '''Feature Filter of EmbeddingVariable Feature'''
                        if args.ev_filter == 'cbf':
                            # CBF-based feature filter
                            filter_option = tf.CBFFilter(
                                filter_freq=3,
                                max_element_size=2**30,
                                false_positive_probability=0.01,
                                counter_type=tf.int64)
                        elif args.ev_filter == 'counter':
                            # Counter-based feature filter
                            filter_option = tf.CounterFilter(filter_freq=3)
                        else:
                            filter_option = None
                        ev_opt = tf.EmbeddingVariableOption(
                            evict_option=evict_opt, filter_option=filter_option)

                        if args.ev:
                            '''Embedding Variable Feature'''
                            categorical_column = tf.feature_column.categorical_column_with_embedding(
                                column_name, dtype=tf.string, ev_option=ev_opt)
                        elif args.adaptive_emb:
                            '''                 Adaptive Embedding Feature Part 2 of 2
                            Expcet the follow code, a dict, 'adaptive_mask_tensors', is need as the input of
                            'tf.feature_column.input_layer(adaptive_mask_tensors=adaptive_mask_tensors)'.
                            For column 'COL_NAME',the value of adaptive_mask_tensors['$COL_NAME'] is a int32
                            tensor with shape [batch_size].
                            '''
                            categorical_column = tf.feature_column.categorical_column_with_adaptive_embedding(
                                column_name,
                                hash_bucket_size=HASH_BUCKET_SIZES[column_name],
                                dtype=tf.string,
                                ev_option=ev_opt)
                        elif args.dynamic_ev:
                            '''Dynamic-dimension Embedding Variable'''
                            print(
                                "Dynamic-dimension Embedding Variable is not really enabled in model."
                            )
                            sys.exit()

                    if args.tf or not args.emb_fusion:
                        embedding_column = tf.feature_column.embedding_column(
                            categorical_column,
                            dimension=64,
                            combiner='mean')
                    else:
                        '''Embedding Fusion Feature'''
                        embedding_column = tf.feature_column.embedding_column(
                            categorical_column,
                            dimension=64,
                            combiner='mean',
                            do_fusion=args.emb_fusion)

                    wide_column.append(embedding_column)
                    deep_column.append(embedding_column)
                    fm_column.append(embedding_column)
                else:
                    column = tf.feature_column.numeric_column(column_name, shape=(1, ))
                    wide_column.append(column)
                    deep_column.append(column)
    else:
        for column_name in FEATURE_COLUMNS:
            if column_name in CATEGORICAL_COLUMNS:
                categorical_column = tf.feature_column.categorical_column_with_hash_bucket(
                    column_name,
                    hash_bucket_size=10000,
                    dtype=tf.string)

                if not args.tf:
                    '''Feature Elimination of EmbeddingVariable Feature'''
                    if args.ev_elimination == 'gstep':
                        # Feature elimination based on global steps
                        evict_opt = tf.GlobalStepEvict(steps_to_live=4000)
                    elif args.ev_elimination == 'l2':
                        # Feature elimination based on l2 weight
                        evict_opt = tf.L2WeightEvict(l2_weight_threshold=1.0)
                    else:
                        evict_opt = None
                    '''Feature Filter of EmbeddingVariable Feature'''
                    if args.ev_filter == 'cbf':
                        # CBF-based feature filter
                        filter_option = tf.CBFFilter(
                            filter_freq=3,
                            max_element_size=2**30,
                            false_positive_probability=0.01,
                            counter_type=tf.int64)
                    elif args.ev_filter == 'counter':
                        # Counter-based feature filter
                        filter_option = tf.CounterFilter(filter_freq=3)
                    else:
                        filter_option = None
                    ev_opt = tf.EmbeddingVariableOption(
                        evict_option=evict_opt, filter_option=filter_option)

                    if args.ev:
                        '''Embedding Variable Feature'''
                        categorical_column = tf.feature_column.categorical_column_with_embedding(
                            column_name, dtype=tf.string, ev_option=ev_opt)
                    elif args.adaptive_emb:
                        '''                 Adaptive Embedding Feature Part 2 of 2
                        Expcet the follow code, a dict, 'adaptive_mask_tensors', is need as the input of
                        'tf.feature_column.input_layer(adaptive_mask_tensors=adaptive_mask_tensors)'.
                        For column 'COL_NAME',the value of adaptive_mask_tensors['$COL_NAME'] is a int32
                        tensor with shape [batch_size].
                        '''
                        categorical_column = tf.feature_column.categorical_column_with_adaptive_embedding(
                            column_name,
                            hash_bucket_size=HASH_BUCKET_SIZES[column_name],
                            dtype=tf.string,
                            ev_option=ev_opt)
                    elif args.dynamic_ev:
                        '''Dynamic-dimension Embedding Variable'''
                        print(
                            "Dynamic-dimension Embedding Variable is not really enabled in model."
                        )
                        sys.exit()

                if args.tf or not args.emb_fusion:
                    embedding_column = tf.feature_column.embedding_column(
                        categorical_column,
                        dimension=64,
                        combiner='mean')
                else:
                    '''Embedding Fusion Feature'''
                    embedding_column = tf.feature_column.embedding_column(
                        categorical_column,
                        dimension=64,
                        combiner='mean',
                        do_fusion=args.emb_fusion)

                wide_column.append(embedding_column)
                deep_column.append(embedding_column)
                fm_column.append(embedding_column)
            else:
                column = tf.feature_column.numeric_column(column_name, shape=(1, ))
                wide_column.append(column)
                deep_column.append(column)
    return wide_column, fm_column, deep_column


def main():
    # check dataset
    print('Checking dataset')
    train_file = args.data_location
    test_file = args.data_location
    if args.parquet_dataset and not args.tf:
        train_file += '/train.parquet'
        test_file += '/eval.parquet'
    else:
        train_file += '/train.csv'
        test_file += '/eval.csv'
    if (not os.path.exists(train_file)) or (not os.path.exists(test_file)):
        print("Dataset does not exist in the given data_location.")
        sys.exit()

    # set fixed random seed
    tf.set_random_seed(args.seed)

    # set directory path
    model_dir = os.path.join(args.output_dir,
                             'model_DeepFM_' + str(int(time.time())))
    checkpoint_dir = args.checkpoint if args.checkpoint else model_dir
    print("Saving model checkpoints to " + checkpoint_dir)

    # create feature column
    wide_column, fm_column, deep_column = build_feature_columns()

    def model_fn(features, labels, mode, config):
        # create variable partitioner for distributed training
        num_ps_replicas = len(json.loads(os.environ["TF_CONFIG"])['cluster']['ps'])
        input_layer_partitioner = partitioned_variables.min_max_variable_partitioner(
            max_partitions=num_ps_replicas,
            min_slice_size=args.input_layer_partitioner <<
            20) if args.input_layer_partitioner else None
        dense_layer_partitioner = partitioned_variables.min_max_variable_partitioner(
            max_partitions=num_ps_replicas,
            min_slice_size=args.dense_layer_partitioner <<
            10) if args.dense_layer_partitioner else None
        print("JUNQI: num_ps_replicas is: ", num_ps_replicas)
        input_layer_partitioner = partitioned_variables.fixed_size_partitioner(3)
        dense_layer_partitioner = partitioned_variables.fixed_size_partitioner(3)
        # create model
        model = DeepFM(wide_column=wide_column,
                    fm_column=fm_column,
                    deep_column=deep_column,
                    optimizer_type=args.optimizer,
                    learning_rate=args.learning_rate,
                    bf16=args.bf16,
                    stock_tf=args.tf,
                    adaptive_emb=args.adaptive_emb,
                    inputs=(features, labels),
                    input_layer_partitioner=input_layer_partitioner,
                    dense_layer_partitioner=dense_layer_partitioner)
        
        model.is_training = True
        return tf.estimator.EstimatorSpec(mode=tf.estimator.ModeKeys.TRAIN, loss=model.loss, train_op=model.train_op,
                                          training_chief_hooks=[ft.train.AsyncCacheCKPTRunnerHook()])

    config = tf.estimator.RunConfig().replace(
        session_config=tf.ConfigProto(device_count={'GPU': 0},
                                      device_filters=cluster_utils.get_device_filters(),
                                      log_device_placement=False,
                                      inter_op_parallelism_threads = args.inter,
                                      intra_op_parallelism_threads = args.intra),
        save_checkpoints_steps=2000, log_step_count_steps=100,
        protocol="elastic-grpc",
        save_summary_steps=50)
    
    esti = tf.estimator.Estimator(
        model_fn = model_fn,
        model_dir=checkpoint_dir,
        config=config)

    train_spec = tf.estimator.TrainSpec(
        input_fn=lambda: build_model_input(train_file, num_epochs=10000, batch_size=args.batch_size),
        max_steps=300000, hooks=[des.python.elastic_training_hooks.ElasticTrainingHook(check_scale_steps=500)])

    eval_spec = tf.estimator.EvalSpec(
        input_fn=lambda: build_model_input(test_file, num_epochs=1, batch_size=args.batch_size),
        # exporters = lastest_exporter,
        steps=None,  # evaluate the whole eval file
        start_delay_secs=1800,
        throttle_secs=600)  # evaluate every 10min for wide
    from dynamic_embedding_server.python import estimator
    estimator.train_and_evaluate(esti, train_spec, eval_spec)



def boolean_string(string):
    low_string = string.lower()
    if low_string not in {'false', 'true'}:
        raise ValueError('Not a valid boolean string')
    return low_string == 'true'


# Get parse
def get_arg_parser():
    parser = argparse.ArgumentParser()
    parser.add_argument('--data_location',
                        help='Full path of train data',
                        required=False,
                        default='./data')
    parser.add_argument('--steps',
                        help='set the number of steps on train dataset',
                        type=int,
                        default=0)
    parser.add_argument('--batch_size',
                        help='Batch size to train. Default is 512',
                        type=int,
                        default=2048)
    parser.add_argument('--output_dir',
                        help='Full path to model output directory. \
                            Default to ./result. Covered by --checkpoint. ',
                        required=False,
                        default='./result')
    parser.add_argument('--checkpoint',
                        help='Full path to checkpoints input/output. \
                            Default to ./result/$MODEL_TIMESTAMP',
                        required=False)
    parser.add_argument('--save_steps',
                        help='set the number of steps on saving checkpoints',
                        type=int,
                        default=0)
    parser.add_argument('--seed',
                        help='set the random seed for tensorflow',
                        type=int,
                        default=2021)
    parser.add_argument('--optimizer',
                        type=str,
                        choices=['adam', 'adamasync',
                                 'adagraddecay', 'adagrad'],
                        default='adamasync')
    parser.add_argument('--learning_rate',
                        help='Learning rate for deep model',
                        type=float,
                        default=0.001)
    parser.add_argument('--keep_checkpoint_max',
                        help='Maximum number of recent checkpoint to keep',
                        type=int,
                        default=1)
    parser.add_argument('--timeline',
                        help='number of steps on saving timeline. Default 0',
                        type=int,
                        default=0)
    parser.add_argument('--protocol',
                        type=str,
                        choices=['grpc', 'grpc++', 'star_server'],
                        default='grpc')
    parser.add_argument('--inter',
                        help='set inter op parallelism threads.',
                        type=int,
                        default=0)
    parser.add_argument('--intra',
                        help='set inter op parallelism threads.',
                        type=int,
                        default=0)
    parser.add_argument('--input_layer_partitioner',
                        help='slice size of input layer partitioner, units MB. Default 8MB',
                        type=int,
                        default=8)
    parser.add_argument('--dense_layer_partitioner',
                        help='slice size of dense layer partitioner, units KB. Default 16KB',
                        type=int,
                        default=16)
    parser.add_argument('--bf16',
                        help='enable DeepRec BF16 in deep model. Default FP32',
                        action='store_true')
    parser.add_argument('--no_eval',
                        help='not evaluate trained model by eval dataset.',
                        action='store_true')
    parser.add_argument('--tf',
                        help='Use TF 1.15.5 API and disable DeepRec feature to run a baseline.',
                        action='store_true')
    parser.add_argument('--smartstaged',
                        help='Whether to enable smart staged feature of DeepRec, Default to True.',
                        type=boolean_string,
                        default=False)
    parser.add_argument('--emb_fusion',
                        help='Whether to enable embedding fusion, Default to True.',
                        type=boolean_string,
                        default=False)
    parser.add_argument('--ev',
                        help='Whether to enable DeepRec EmbeddingVariable. Default False.',
                        type=boolean_string,
                        default=False)
    parser.add_argument('--ev_elimination',
                        help='Feature Elimination of EmbeddingVariable Feature. Default closed.',
                        type=str,
                        choices=[None, 'l2', 'gstep'],
                        default=None)
    parser.add_argument('--ev_filter',
                        help='Feature Filter of EmbeddingVariable Feature. Default closed.',
                        type=str,
                        choices=[None, 'counter', 'cbf'],
                        default=None)
    parser.add_argument('--op_fusion',
                        help='Whether to enable Auto graph fusion feature. Default to True',
                        type=boolean_string,
                        default=True)
    parser.add_argument('--micro_batch',
                        help='Set num for Auto Mirco Batch. Default close.',
                        type=int,
                        default=0) # TODO enable
    parser.add_argument('--adaptive_emb',
                        help='Whether to enable Adaptive Embedding. Default to False.',
                        type=boolean_string,
                        default=False)
    parser.add_argument('--dynamic_ev',
                        help='Whether to enable Dynamic-dimension Embedding Variable. Default to False.',
                        type=boolean_string,
                        default=False) # TODO enable
    parser.add_argument('--incremental_ckpt',
                        help='Set time of save Incremental Checkpoint. Default 0 to close.',
                        type=boolean_string,
                        default=0)
    parser.add_argument('--workqueue',
                        help='Whether to enable Work Queue. Default to False.',
                        type=boolean_string,
                        default=False)
    parser.add_argument("--parquet_dataset", \
                        help='Whether to enable Parquet DataSet. Defualt to True.',
                        type=boolean_string,
                        default=False)
    parser.add_argument("--parquet_dataset_shuffle", \
                        help='Whether to enable shuffle operation for Parquet Dataset. Default to False.',
                        type=boolean_string,
                        default=False)
    parser.add_argument("--group_embedding", \
                        help='Whether to enable Group Embedding. Defualt to None.',
                        type=str,
                        choices=[None, 'localized', 'collective'],
                        default=None)

    return parser


# Some DeepRec's features are enabled by ENV.
# This func is used to set ENV and enable these features.
# A triple quotes comment is used to introduce these features and play an emphasizing role.
def set_env_for_DeepRec():
    '''
    Set some ENV for these DeepRec's features enabled by ENV.
    More Detail information is shown in https://deeprec.readthedocs.io/zh/latest/index.html.
    START_STATISTIC_STEP & STOP_STATISTIC_STEP: On CPU platform, DeepRec supports memory optimization
        in both stand-alone and distributed trainging. It's default to open, and the
        default start and stop steps of collection is 1000 and 1100. Reduce the initial
        cold start time by the following settings.
    MALLOC_CONF: On CPU platform, DeepRec can use memory optimization with the jemalloc library.
        Please preload libjemalloc.so by `LD_PRELOAD=./libjemalloc.so.2 python ...`
    '''
    os.environ['START_STATISTIC_STEP'] = '100'
    os.environ['STOP_STATISTIC_STEP'] = '110'
    os.environ['MALLOC_CONF'] = \
        'background_thread:true,metadata_thp:auto,dirty_decay_ms:20000,muzzy_decay_ms:20000'
    if args.group_embedding == "collective":
        tf.config.experimental.enable_distributed_strategy(strategy="collective")
        if args.smartstaged and not args.tf:
            os.environ["TF_GPU_THREAD_COUNT"] = "16"

if __name__ == '__main__':
    parser = get_arg_parser()
    args = parser.parse_args()

    set_env_for_DeepRec()
    main()