如何实现支持多值、稀疏、共享权重的 DeepFM

def _decode_tsv(line):    columns = tf.decode_csv(line, record_defaults=DEFAULT_VALUES, field_delim='\t')    y = columns[0]
    feat_columns = dict(zip((t[0] for t in COLUMNS_MAX_TOKENS), columns[1:]))    X = {}    for colname, max_tokens in COLUMNS_MAX_TOKENS:        # 调用string_split时，第一个参数必须是一个list，所以要把columns[colname]放在[]中        # 这时每个kv还是'k:v'这样的字符串        kvpairs = tf.string_split([feat_columns[colname]], ',').values[:max_tokens]
        # k,v已经拆开, kvpairs是一个SparseTensor，因为每个kvpair格式相同，都是"k:v"        # 既不会出现"k"，也不会出现"k:v1:v2:v3:..."        # 所以，这时的kvpairs实际上是一个满阵        kvpairs = tf.string_split(kvpairs, ':')
        # kvpairs是一个[n_valid_pairs,2]矩阵        kvpairs = tf.reshape(kvpairs.values, kvpairs.dense_shape)
        feat_ids, feat_vals = tf.split(kvpairs, num_or_size_splits=2, axis=1)        feat_ids = tf.string_to_number(feat_ids, out_type=tf.int32)        feat_vals = tf.string_to_number(feat_vals, out_type=tf.float32)
        # 不能调用squeeze, squeeze的限制太多, 当原始矩阵有1行或0行时，squeeze都会报错        X[colname + "_ids"] = tf.reshape(feat_ids, shape=[-1])        X[colname + "_values"] = tf.reshape(feat_vals, shape=[-1])
    return X, y

def input_fn(data_file, n_repeat, batch_size, batches_per_shuffle):    # ----------- prepare padding    pad_shapes = {}    pad_values = {}    for c, max_tokens in COLUMNS_MAX_TOKENS:        pad_shapes[c + "_ids"] = tf.TensorShape([max_tokens])        pad_shapes[c + "_values"] = tf.TensorShape([max_tokens])
        pad_values[c + "_ids"] = -1  # 0 is still valid token-id, -1 for padding        pad_values[c + "_values"] = 0.0
    # no need to pad labels    pad_shapes = (pad_shapes, tf.TensorShape([]))    pad_values = (pad_values, 0)
    # ----------- define reading ops    dataset = tf.data.TextLineDataset(data_file).skip(1)  # skip the header    dataset = dataset.map(_decode_tsv, num_parallel_calls=4)
    if batches_per_shuffle > 0:        dataset = dataset.shuffle(batches_per_shuffle * batch_size)
    dataset = dataset.repeat(n_repeat)    dataset = dataset.padded_batch(batch_size=batch_size,                                   padded_shapes=pad_shapes,                                   padding_values=pad_values)
    iterator = dataset.make_one_shot_iterator()    dense_Xs, ys = iterator.get_next()
    # ----------- convert dense to sparse    sparse_Xs = {}    for c, _ in COLUMNS_MAX_TOKENS:        for suffix in ["ids", "values"]:            k = "{}_{}".format(c, suffix)            sparse_Xs[k] = tf_utils.to_sparse_input_and_drop_ignore_values(dense_Xs[k])
    # ----------- return    return sparse_Xs, ys

class EmbeddingTable:    def __init__(self):        self._weights = {}
    def add_weights(self, vocab_name, vocab_size, embed_dim):        """        :param vocab_name: 一个field拥有两个权重矩阵，一个用于线性连接，另一个用于非线性（二阶或更高阶交叉）连接        :param vocab_size: 字典总长度        :param embed_dim: 二阶权重矩阵shape=[vocab_size, order2dim]，映射成的embedding                          既用于接入DNN的第一屋，也是用于FM二阶交互的隐向量        :return: None        """        linear_weight = tf.get_variable(name='{}_linear_weight'.format(vocab_name),                                        shape=[vocab_size, 1],                                        initializer=tf.glorot_normal_initializer(),                                        dtype=tf.float32)
        # 二阶（FM）与高阶（DNN）的特征交互，共享embedding矩阵        embed_weight = tf.get_variable(name='{}_embed_weight'.format(vocab_name),                                       shape=[vocab_size, embed_dim],                                       initializer=tf.glorot_normal_initializer(),                                       dtype=tf.float32)
        self._weights[vocab_name] = (linear_weight, embed_weight)
    def get_linear_weights(self, vocab_name): return self._weights[vocab_name][0]
    def get_embed_weights(self, vocab_name): return self._weights[vocab_name][1]

def build_embedding_table(params):    embed_dim = params['embed_dim']  # 必须有统一的embedding长度
    embedding_table = EmbeddingTable()    for vocab_name, vocab_size in params['vocab_sizes'].items():        embedding_table.add_weights(vocab_name=vocab_name, vocab_size=vocab_size, embed_dim=embed_dim)
    return embedding_table

def output_logits_from_linear(features, embedding_table, params):    field2vocab_mapping = params['field_vocab_mapping']    combiner = params.get('multi_embed_combiner', 'sum')
    fields_outputs = []    # 当前field下有一系列的<tag:value>对，每个tag对应一个bias（待优化），    # 将所有tag对应的bias，按照其value进行加权平均，得到这个field对应的bias    for fieldname, vocabname in field2vocab_mapping.items():        sp_ids = features[fieldname + "_ids"]        sp_values = features[fieldname + "_values"]
        linear_weights = embedding_table.get_linear_weights(vocab_name=vocabname)
        # weights: [vocab_size,1]        # sp_ids: [batch_size, max_tags_per_example]        # sp_weights: [batch_size, max_tags_per_example]        # output: [batch_size, 1]        output = embedding_ops.safe_embedding_lookup_sparse(linear_weights, sp_ids, sp_values,                                                            combiner=combiner,                                                            name='{}_linear_output'.format(fieldname))
        fields_outputs.append(output)
    # 因为不同field可以共享同一个vocab的linear weight，所以将各个field的output相加，会损失大量的信息    # 因此，所有field对应的output拼接起来，反正每个field的output都是[batch_size,1]，拼接起来，并不占多少空间    # whole_linear_output: [batch_size, total_fields]    whole_linear_output = tf.concat(fields_outputs, axis=1)    tf.logging.info("linear output, shape={}".format(whole_linear_output.shape))
    # 再映射到final logits（二分类，也是[batch_size,1]）    # 这时，就不要用任何activation了，特别是ReLU    return tf.layers.dense(whole_linear_output, units=1, use_bias=True, activation=None)

def output_logits_from_bi_interaction(features, embedding_table, params):    field2vocab_mapping = params['field_vocab_mapping']
    # 论文上的公式就是要求sum，而且我也试过mean和sqrtn，都比用mean要差上很多    # 但是，这种情况，仅仅是针对criteo数据的，还是理论上就必须用sum，而不能用mean和sqrtn    # 我还不太确定，所以保留一个接口能指定其他combiner的方法    combiner = params.get('multi_embed_combiner', 'sum')
    # 见《Neural Factorization Machines for Sparse Predictive Analytics》论文的公式(4)    fields_embeddings = []    fields_squared_embeddings = []
    for fieldname, vocabname in field2vocab_mapping.items():        sp_ids = features[fieldname + "_ids"]        sp_values = features[fieldname + "_values"]
        # --------- embedding        embed_weights = embedding_table.get_embed_weights(vocabname)        # embedding: [batch_size, embed_dim]        embedding = embedding_ops.safe_embedding_lookup_sparse(embed_weights, sp_ids, sp_values,                                                               combiner=combiner,                                                               name='{}_embedding'.format(fieldname))        fields_embeddings.append(embedding)
        # --------- square of embedding        squared_emb_weights = tf.square(embed_weights)
        squared_sp_values = tf.SparseTensor(indices=sp_values.indices,                                            values=tf.square(sp_values.values),                                            dense_shape=sp_values.dense_shape)
        # squared_embedding: [batch_size, embed_dim]        squared_embedding = embedding_ops.safe_embedding_lookup_sparse(squared_emb_weights, sp_ids, squared_sp_values,                                                                       combiner=combiner,                                                                       name='{}_squared_embedding'.format(fieldname))        fields_squared_embeddings.append(squared_embedding)
    # calculate bi-interaction    sum_embedding_then_square = tf.square(tf.add_n(fields_embeddings))  # [batch_size, embed_dim]    square_embedding_then_sum = tf.add_n(fields_squared_embeddings)  # [batch_size, embed_dim]    bi_interaction = 0.5 * (sum_embedding_then_square - square_embedding_then_sum)  # [batch_size, embed_dim]    tf.logging.info("bi-interaction, shape={}".format(bi_interaction.shape))
    # calculate logits    logits = tf.layers.dense(bi_interaction, units=1, use_bias=True, activation=None)
    # 因为FM与DNN共享embedding，所以除了logits，还返回各field的embedding，方便搭建DNN    return logits, fields_embeddings

def output_logits_from_dnn(fields_embeddings, params, is_training):    dropout_rate = params['dropout_rate']    do_batch_norm = params['batch_norm']
    X = tf.concat(fields_embeddings, axis=1)    tf.logging.info("initial input to DNN, shape={}".format(X.shape))
    for idx, n_units in enumerate(params['hidden_units'], start=1):        X = tf.layers.dense(X, units=n_units, activation=tf.nn.relu)        tf.logging.info("layer[{}] output shape={}".format(idx, X.shape))
        X = tf.layers.dropout(inputs=X, rate=dropout_rate, training=is_training)        if is_training:            tf.logging.info("layer[{}] dropout {}".format(idx, dropout_rate))
        if do_batch_norm:            # BatchNormalization的调用、参数，是从DNNLinearCombinedClassifier源码中拷贝过来的            batch_norm_layer = normalization.BatchNormalization(momentum=0.999, trainable=True,                                                                name='batchnorm_{}'.format(idx))            X = batch_norm_layer(X, training=is_training)
            if is_training:                tf.logging.info("layer[{}] batch-normalize".format(idx))
    # connect to final logits, [batch_size,1]    return tf.layers.dense(X, units=1, use_bias=True, activation=None)

def model_fn(features, labels, mode, params):    for featname, featvalues in features.items():        if not isinstance(featvalues, tf.SparseTensor):            raise TypeError("feature[{}] isn't SparseTensor".format(featname))
    # ============= build the graph    embedding_table = build_embedding_table(params)
    linear_logits = output_logits_from_linear(features, embedding_table, params)
    bi_interact_logits, fields_embeddings = output_logits_from_bi_interaction(features, embedding_table, params)
    dnn_logits = output_logits_from_dnn(fields_embeddings, params, (mode == tf.estimator.ModeKeys.TRAIN))
    general_bias = tf.get_variable(name='general_bias', shape=[1], initializer=tf.constant_initializer(0.0))
    logits = linear_logits + bi_interact_logits + dnn_logits    logits = tf.nn.bias_add(logits, general_bias)  # bias_add，获取broadcasting的便利
    # reshape [batch_size,1] to [batch_size], to match the shape of 'labels'    logits = tf.reshape(logits, shape=[-1])
    probabilities = tf.sigmoid(logits)
    # ============= predict spec    if mode == tf.estimator.ModeKeys.PREDICT:        return tf.estimator.EstimatorSpec(            mode=mode,            predictions={'probabilities': probabilities})
    # ============= evaluate spec    # STUPID TENSORFLOW CANNOT AUTO-CAST THE LABELS FOR ME    loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=logits, labels=tf.cast(labels, tf.float32)))
    eval_metric_ops = {'auc': tf.metrics.auc(labels, probabilities)}    if mode == tf.estimator.ModeKeys.EVAL:        return tf.estimator.EstimatorSpec(            mode=mode,            loss=loss,            eval_metric_ops=eval_metric_ops)
    # ============= train spec    assert mode == tf.estimator.ModeKeys.TRAIN    train_op = params['optimizer'].minimize(loss, global_step=tf.train.get_global_step())    return tf.estimator.EstimatorSpec(mode,                                      loss=loss,                                      train_op=train_op,                                      eval_metric_ops=eval_metric_ops)

def get_hparams():    vocab_sizes = {        'numeric': 13,        # there are totally 14738 categorical tags occur >= 200        # since 0 is reserved for OOV, so total vocab_size=14739        'categorical': 14739    }
    optimizer = tf.train.ProximalAdagradOptimizer(        learning_rate=0.01,        l1_regularization_strength=0.001,        l2_regularization_strength=0.001)
    return {        'embed_dim': 128,        'vocab_sizes': vocab_sizes,        # 在这个case中，没有多个field共享同一个vocab的情况，而且field_name和vocab_name相同        'field_vocab_mapping': {'numeric': 'numeric', 'categorical': 'categorical'},        'dropout_rate': 0.3,        'batch_norm': False,        'hidden_units': [64, 32],        'optimizer': optimizer    }

if __name__ == "__main__":    tf.logging.set_verbosity(tf.logging.INFO)    tf.set_random_seed(999)
    hparams = get_hparams()    deepfm = tf.estimator.Estimator(model_fn=model_fn,                                    model_dir='models/criteo',                                    params=hparams)
    train_spec = tf.estimator.TrainSpec(input_fn=lambda: input_fn(data_file='dataset/criteo/whole_train.tsv',                                                                  n_repeat=10,                                                                  batch_size=128,                                                                  batches_per_shuffle=10))
    eval_spec = tf.estimator.EvalSpec(input_fn=lambda: input_fn(data_file='dataset/criteo/whole_test.tsv',                                                                n_repeat=1,                                                                batch_size=128,                                                                batches_per_shuffle=-1))
    tf.estimator.train_and_evaluate(deepfm, train_spec, eval_spec)