利用小trick加速tensorflow的训练

tensorflow ==1.13.1

1. tf.data并行读取tfrecord数据

def parse_exp(example):
    features = {}
    """ tfrecord解析代码 """
    return features

def input_fn(filenames = "./train_data/*.tfrecord", batch_size=128):
    files = tf.data.Dataset.list_files(filenames)
    dataset = files.apply(tf.contrib.data.parallel_interleave(lambda filename:         
    tf.data.TFRecordDataset(files), buffer_output_elements=batch_size*20, cycle_length=10))
    dataset = dataset.shuffle(batch_size*4)
    dataset = dataset.map(parse_exp, num_parallel_calls=8)
    dataset = dataset.repeat().batch(batch_size).prefetch(1)
    return dataset
    

2. 在进行分布式训练时，使用tf.fixed_size_partitioner参数分割，对于有较大Embedding计算的时候尤其有用，代码如下：

def model_fn(features, mode, params):
    """ 构建estimator模型 """
    with tf.variable_scope("deviceID_embedding", partitioner=tf.fixed_size_partitioner(8, axis=0)):
        deviceID_input = tf.feature_column.input_layer(features, params["deviceID"])

    """ 构建自己的代码逻辑 """
    net = ...
    output = tf.layers.dense(net, units=1)
    return output
    

   其中tf.fixed_size_partitioner(8, axis=0)的 8代表ps个数。