欢迎您访问程序员文章站本站旨在为大家提供分享程序员计算机编程知识!
您现在的位置是: 首页

利用小trick加速tensorflow的训练

程序员文章站 2024-03-23 13:46:28
...

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(32, axis=0)):
        deviceID_input = tf.feature_column.input_layer(features, params["deviceID"])

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