wide&deep模型 迁移样例（基于DLRM模型框架）

https://github.com/ZiyaoGeng/RecLearn

Criteo4500w数据集:  
https://ailab.criteo.com/ressources/kaggle-display-advertising-challenge-dataset.tar.gz

python critro.py --data_path data_path --output_path output_path 

# get txt_list
file_split_list = get_split_file_path(dataset_path=data_path)

# get feature_map
feature_map = get_fea_map(split_file_list=file_split_list)

# dense feature: Bin continuous data into intervals.
data_df[dense_features] = rec_kbins_discretizer(data_df[dense_features], 1000, feature_map)

# sparse feature: mapping
for col in sparse_features:
    try:
        data_df[col] = data_df[col].map(lambda x: feature_map[col][x])
    except KeyError as e:
        raise KeyError("Feature {} not found in dataset".format(col)) from e

# add offsets
slot_size_array = [
                1001, 1001, 1001, 1001, 1001, 1001, 1001, 1001, 1001, 1001, 1001, 1001, 1001,
                1462, 585, 10131228, 2202609, 307, 25, 12519, 635, 5, 93147, 5685, 8351594, 3196,
                29, 14994, 5461307, 12, 5654, 2174, 5, 7046548, 19, 17, 286182, 106, 142573
]
offset_size_list = np.cumsum([0] + slot_size_array[:-1])
for col_index in range(1, len(offset_size_list) + 1):
    data_df.iloc[:, col_index] += offset_size_list[col_index - 1]

# txt to tfrecords
convert_input2tfrd(in_file_path=file, out_file_path=output_path)

bash run.sh {so_path} {rec_package_path} {hccl_cfg_json} {dlrm_criteo_data_path}

bash run.sh {so_path} {rec_package_path} {hccl_cfg_json} {dlrm_criteo_data_path} {IP}

# run.sh: 32~37行
export USE_DYNAMIC=0            # 0：静态shape；1：动态shape
export CACHE_MODE="HBM"         # HBM；DDR；SSD
export USE_FAAE=0               # 0：关闭准入淘汰；1：开启准入淘汰
export USE_DYNAMIC_EXPANSION=0  # 0：关闭动态扩容；1: 开启动态扩容
export USE_MULTI_LOOKUP=0       # 0：一表一查；1：一表多查
export USE_MODIFY_GRAPH=0       # 0：feature spec模式；1：自动改图模式

# 88~89行
lr_sparse = self.base_lr_sparse * lr_factor_constant
lr_dense = self.base_lr_dense * lr_factor_constant
# 140~146行
_lr_scheduler = LearningRateScheduler(
    0.001,
    0.001,
    LR_SCHEDULE_STEPS[0],
    LR_SCHEDULE_STEPS[1],
    LR_SCHEDULE_STEPS[2],
)
# 超参数
self.batch_size = 4096
self.line_per_sample = 1
self.train_epoch = 1
self.test_epoch = 9
self.emb_dim = 8           

# wdl.py
import tensorflow as tf
from tensorflow.keras import Model
from tensorflow.keras.layers import Dense, Embedding, Dropout, Input
from tensorflow.keras.regularizers import l2

from reclearn.layers import Linear, MLP
from reclearn.layers.utils import index_mapping

class WideDeep(Model):
    def __init__(self, feature_columns, hidden_units, activation='relu',
                 dnn_dropout=0., embed_reg=0., w_reg=0.):
        """Wide&Deep.
        Args:
            :param feature_columns: A list. [{'feat_name':, 'feat_num':, 'embed_dim':}, ...]
            :param hidden_units: A list. Neural network hidden units.
            :param activation: A string. Activation function of MLP.
            :param dnn_dropout: A scalar. Dropout of MLP.
            :param embed_reg: A scalar. The regularization coefficient of embedding.
            :param w_reg: A scalar. The regularization coefficient of Linear.
        :return
        """
        super(WideDeep, self).__init__()
        self.feature_columns = feature_columns
        self.embed_layers = {
            feat['feat_name']: Embedding(input_dim=feat['feat_num'],
                                         input_length=1,
                                         output_dim=feat['embed_dim'],
                                         embeddings_initializer='random_normal',
                                         embeddings_regularizer=l2(embed_reg))
            for feat in self.feature_columns
        }
        self.map_dict = {}
        self.feature_length = 0
        for feat in self.feature_columns:
            self.map_dict[feat['feat_name']] = self.feature_length
            self.feature_length += feat['feat_num']
        self.dnn_network = MLP(hidden_units, activation, dnn_dropout)
        self.linear = Linear(self.feature_length, w_reg=w_reg)
        self.final_dense = Dense(1, activation=None)

    def call(self, inputs):
        sparse_embed = tf.concat([self.embed_layers[feat_name](value) for feat_name, value in inputs.items()], axis=-1)
        x = sparse_embed  # (batch_size, field * embed_dim)
        # Wide
        wide_inputs = index_mapping(inputs, self.map_dict)
        wide_inputs = tf.concat([value for _, value in wide_inputs.items()], axis=-1)
        wide_out = self.linear(wide_inputs)
        # Deep
        deep_out = self.dnn_network(x)
        deep_out = self.final_dense(deep_out)
        # out
        outputs = tf.nn.sigmoid(0.5 * wide_out + 0.5 * deep_out)
        return outputs

    def summary(self):
        inputs = {
            feat['feat_name']: Input(shape=(), dtype=tf.int32, name=feat['feat_name'])
            for feat in self.feature_columns
        }
        Model(inputs=inputs, outputs=self.call(inputs)).summary()

# model.py
import time
from easydict import EasyDict as edict

import tensorflow as tf


model_cfg = edict()
model_cfg.loss_mode = "batch"
LOSS_OP_NAME = "loss"
LABEL_OP_NAME = "label"
VAR_LIST = "variable"
PRED_OP_NAME = "pred"


class MyModel:
    def __init__(self):
        self.kernel_init = None
        self._loss_fn = None
        self.is_training = None

    def build_model(self,
                    wide_embedding=None,
                    deep_embedding=None,
                    label=None,
                    is_training=True,
                    seed=None,
                    dropout_rate=None,
                    batch_norm=False):

        with tf.variable_scope("wide_deep", reuse=tf.AUTO_REUSE):
            self._loss_fn = tf.keras.losses.BinaryCrossentropy(from_logits=True)
            self.is_training = is_training

            # wide
            batch_size, wide_num, wide_emb_dim = wide_embedding.shape
            wide_input = tf.reshape(wide_embedding[:,0], shape=(batch_size, wide_num * 1))
            wide_output = tf.reshape(tf.reduce_sum(wide_input, axis=1), shape=(-1,1))

            # deep
            batch_size, deep_num, deep_emb_dim = deep_embedding.shape
            deep_input = tf.reshape(deep_embedding, shape=(batch_size, deep_num * deep_emb_dim))

            ## MLP
            hidden_units = [256,128,64]
            net = deep_input
            for i,unit in enumerate(hidden_units):

                net = tf.layers.dense(net, units=unit, activation='relu', name=f'hidden_layer_{i}',
                                      kernel_initializer=tf.glorot_uniform_initializer(seed=seed),
                                      bias_initializer=tf.zeros_initializer())

                if dropout_rate is not None and 0.0 < dropout_rate < 1.0:
                    net = tf.layers.dropout(net,dropout_rate,training=self.is_training)
                if batch_norm:
                    net = tf.layers.batch_normalization(net, training=self.is_training)
                    
            deep_output = tf.layers.dense(net, units=1, activation=None, name='deep_output',
                                          kernel_initializer=tf.glorot_uniform_initializer(seed=seed),
                                          bias_initializer=tf.zeros_initializer())
            
            total_logits = 0.5 * tf.add(wide_output,deep_output,name='total_logits')
            loss = self._loss_fn(label, total_logits)
            prediction = tf.sigmoid(total_logits)
            trainable_variables = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope='wide_deep')
            return {LOSS_OP_NAME: loss,
                    PRED_OP_NAME: prediction,
                    LABEL_OP_NAME: label,
                    VAR_LIST: trainable_variables}


my_model = MyModel()

# main_mxrec.py：100~104行
def map_fn(batch):
    new_batch = batch
    new_batch['sparse_feature'] = tf.concat([batch['dense_feature'], batch['sparse_feature']], axis=1)
    return new_batch
dataset = dataset.map(map_fn, num_parallel_calls=num_parallel)

# main_mxrec.py：109~118行
if not MODIFY_GRAPH_FLAG:

    # Enable EOSDataset manually.
    librec = import_host_pipeline_ops(LIBREC_EOS_OPS_SO)
    channel_id = 0 if is_training else 1
    # 此处eos_map的调用必须先于insert_func,避免多卡数据不均匀的情况
    dataset = dataset.eos_map(librec, channel_id, kwargs.get("max_train_steps", max_train_steps),
                              kwargs.get("max_eval_steps", eval_steps))
    insert_fn = get_asc_insert_func(tgt_key_specs=feature_spec_list, is_training=is_training, dump_graph=dump_graph)
    dataset = dataset.map(insert_fn)

# main_mxrec.py：127~179行
def model_forward(feature_list, wide_hash_table_list, deep_hash_table_list, batch, is_train, modify_graph, is_use_faae=False):
    wide_embedding_list = []
    deep_embedding_list = []
    wide_feature_list = []
    deep_feature_list = []
    if is_use_faae:
        feature_list_copy = feature_list[:-1]
    else:
        feature_list_copy = feature_list

    for i,item in enumerate(feature_list_copy):
        if i % 2 == 0:
            wide_feature_list.append(item)
        else:
            deep_feature_list.append(item)

    logger.debug(f"In model_forward function, is_train: {is_train}, feature_list: {len(feature_list)}, "
                 f"wide_hash_table_list: {len(wide_hash_table_list)}, deep_hash_table_list: {len(deep_hash_table_list)}")

    # wide
    for wide_feature, wide_hash_table in zip(wide_feature_list, wide_hash_table_list):
        if MODIFY_GRAPH_FLAG:
            wide_feature = batch["sparse_feature"]
        wide_embedding = sparse_lookup(wide_hash_table, wide_feature, cfg.send_count, dim=None, is_train=is_train,
                                  name="wide_embedding_lookup", modify_graph=modify_graph, batch=batch,
                                  access_and_evict_config=None)
        wide_embedding_list.append(wide_embedding)

    # deep
    for deep_feature, deep_hash_table in zip(deep_feature_list, deep_hash_table_list):
        if MODIFY_GRAPH_FLAG:
            deep_feature = batch["sparse_feature"]
        deep_embedding = sparse_lookup(deep_hash_table, deep_feature, cfg.send_count, dim=None, is_train=is_train,
                                  name="deep_embedding_lookup", modify_graph=modify_graph, batch=batch,
                                  access_and_evict_config=None)
        deep_embedding_list.append(deep_embedding)

    if len(wide_embedding_list) == 1:
        wide_emb = wide_embedding_list[0]
        deep_emb = deep_embedding_list[0]
    elif len(wide_embedding_list) > 1:
        wide_emb = tf.reduce_sum(wide_embedding_list, axis=0, keepdims=False)
        deep_emb = tf.reduce_sum(deep_embedding_list, axis=0, keepdims=False)
    else:
        raise ValueError("the length of embedding_list must be greater than or equal to 1.")
    my_model = MyModel()
    model_output = my_model.build_model(wide_embedding=wide_emb,
                                        deep_embedding=deep_emb,
                                        label=batch["label"],
                                        is_training=is_train,
                                        seed=dense_hashtable_seed,
                                        dropout_rate=0.5)
    return model_output

# main_mxrec.py: 273~296行
def create_feature_spec_list(use_timestamp=False):
    access_threshold = None
    eviction_threshold = None
    if use_timestamp:
        access_threshold = 1000
        eviction_threshold = 180

    feature_spec_list = [FeatureSpec("sparse_feature", table_name="wide_embeddings", batch_size=cfg.batch_size,
                                     access_threshold=access_threshold, eviction_threshold=eviction_threshold),
                         FeatureSpec("sparse_feature", table_name="deep_embeddings", batch_size=cfg.batch_size,
                                     access_threshold=access_threshold, eviction_threshold=eviction_threshold)]

    if use_multi_lookup:
        feature_spec_list.extend([FeatureSpec("sparse_feature", table_name="wide_embeddings",
                                             batch_size=cfg.batch_size,
                                             access_threshold=access_threshold,
                                             eviction_threshold=eviction_threshold),
                                  FeatureSpec("sparse_feature", table_name="deep_embeddings",
                                             batch_size=cfg.batch_size,
                                             access_threshold=access_threshold,
                                             eviction_threshold=eviction_threshold)])
    if use_timestamp:
        feature_spec_list.append(FeatureSpec("timestamp", is_timestamp=True))
    return feature_spec_list

# main_mxrec.py: 379~397行
# 创表操作
wide_emb_initializer = tf.compat.v1.truncated_normal_initializer(stddev=0.05, seed=sparse_hashtable_seed)
deep_emb_initializer = tf.compat.v1.truncated_normal_initializer(stddev=0.05, seed=sparse_hashtable_seed)

sparse_hashtable_wide = create_table(
    key_dtype=cfg.key_type,
    dim=tf.TensorShape([cfg.emb_dim]),
    name="wide_embeddings",
    emb_initializer=wide_emb_initializer,
    **cfg.get_emb_table_cfg()
)

sparse_hashtable_deep = create_table(
    key_dtype=cfg.key_type,
    dim=tf.TensorShape([cfg.emb_dim]),
    name="deep_embeddings",
    emb_initializer=deep_emb_initializer,
    **cfg.get_emb_table_cfg()
)

# main_mxrec.py: 410~442行
train_variables, emb_variables = get_dense_and_sparse_variable()

rank_size = mxrec_util.communication.hccl_ops.get_rank_size()
train_ops = []
# multi task training
for loss, (model_optimizer, emb_optimizer) in zip([train_model.get("loss")], optimizer_list):
    # do model optimization
    grads = model_optimizer.compute_gradients(loss, var_list=train_variables)
    avg_grads = []
    for grad, var in grads:
        if rank_size > 1:
            grad = hccl_ops.allreduce(grad, "sum") if grad is not None else None
        if grad is not None:
            avg_grads.append((grad / 8.0, var))
    # apply gradients: update variables
    train_ops.append(model_optimizer.apply_gradients(avg_grads))

    if use_dynamic_expansion:
        train_address_list = tf.compat.v1.get_collection(ASCEND_SPARSE_LOOKUP_ID_OFFSET)
        train_emb_list = tf.compat.v1.get_collection(ASCEND_SPARSE_LOOKUP_LOCAL_EMB)
        # do embedding optimization by addr
        sparse_grads = emb_optimizer.compute_gradients(loss, train_emb_list)  # local_embedding
        grads_and_vars = [(grad, address) for grad, address in zip(sparse_grads, train_address_list)]
        train_ops.append(emb_optimizer.apply_gradients(grads_and_vars))
    else:
        # do embedding optimization
        sparse_grads = emb_optimizer.compute_gradients(loss, emb_variables)
        print("sparse_grads_tensor:", sparse_grads)
        grads_and_vars = [(grad, variable) for grad, variable in zip(sparse_grads, emb_variables)]
        train_ops.append(emb_optimizer.apply_gradients(grads_and_vars))

# 动态学习率更新
train_ops.extend([cfg.global_step.assign(cfg.global_step + 1), cfg.learning_rate[0], cfg.learning_rate[1]])

import tensorflow as tf
from delay_loss_scale import DenseLossScaleOptimizer, SparseLossScaleOptimizer
from mx_rec.util.initialize import ConfigInitializer
from mx_rec.optimizers.lazy_adam import create_hash_optimizer
from mx_rec.optimizers.lazy_adam_by_addr import create_hash_optimizer_by_address


def get_dense_and_sparse_optimizer(cfg):
    dense_optimizer = tf.train.AdamOptimizer(learning_rate=cfg.learning_rate[0])
    use_dynamic_expansion = ConfigInitializer.get_instance().use_dynamic_expansion
    if use_dynamic_expansion:
        sparse_optimizer = create_hash_optimizer_by_address(learning_rate=cfg.learning_rate[1])
    else:
        sparse_optimizer = create_hash_optimizer(learning_rate=cfg.learning_rate[1])
    sparse_optimizer = SparseLossScaleOptimizer(sparse_optimizer, 1)
    dense_optimizer = DenseLossScaleOptimizer(dense_optimizer, 1)

    return dense_optimizer, sparse_optimizer