利用小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個數。