利用TF dataset改善模型训练效率的最佳实践

Posted 2022-01-19 MrCharles

不好的实践

已经提前把数据全部保存为tfrecord， 以便提升模型训练的时候的效率，数据集大小大概为4G左右。使用如下数据集构建流程：

def load_tfrecord_dataset(tfrecord_name, batch_size, shuffle=True,
                          buffer_size=1024):
    """load dataset from tfrecord"""
    raw_dataset = tf.data.TFRecordDataset(tfrecord_name)#.cache()
    raw_dataset = raw_dataset.repeat()
    if shuffle:
        raw_dataset = raw_dataset.shuffle(
            buffer_size=buffer_size)
    dataset = raw_dataset.map(
        _parse_tfrecord(),
        num_parallel_calls=tf.data.experimental.AUTOTUNE)
    dataset = dataset.batch(batch_size).map( lambda x, y: randomcrop(x, y),  num_parallel_calls=tf.data.experimental.AUTOTUNE).prefetch(FLAGS.prefetch)
    # dataset = dataset.prefetch(
    #     buffer_size=tf.data.experimental.AUTOTUNE)
    return dataset

很显然，这样是不好的操作，只能达到70左右的利用率：

一个epoch需要34S：

好的实践

修改为如下流程，既可以提升很好：

def load_tfrecord_dataset_good(tfrecord_name, batch_size, shuffle=True,
                          buffer_size=1024):
    """load dataset from tfrecord"""
    raw_dataset = tf.data.TFRecordDataset(tfrecord_name)
    dataset = raw_dataset.map(_parse_tfrecord(),num_parallel_calls=tf.data.experimental.AUTOTUNE).cache().repeat().shuffle(buffer_size=buffer_size)
    dataset = dataset.batch(batch_size).map( lambda x, y: randomcrop(x, y),  num_parallel_calls=tf.data.experimental.AUTOTUNE)
    dataset = dataset.prefetch(buffer_size=tf.data.experimental.AUTOTUNE)
    return dataset

可以提升到86左右了：

一个epoch大概27S：

可是还是没有跑满GPU

更好的做法

注意到我们随机裁剪是在batch之后，可能会受限制与batch的操作，导致来不及处理，我们可以让他提前裁剪，然后再去取batch，这样可以加快，让CPU一直去裁剪：

def load_tfrecord_dataset_better(tfrecord_name, batch_size, shuffle=True,
                          buffer_size=1024):
    """load dataset from tfrecord"""
    raw_dataset = tf.data.TFRecordDataset(tfrecord_name)
    dataset = raw_dataset.map(_parse_tfrecord(),num_parallel_calls=tf.data.experimental.AUTOTUNE).cache().repeat().shuffle(buffer_size=buffer_size)
    dataset = dataset.map( lambda x, y: randomcrop_single(x, y),  num_parallel_calls=tf.data.experimental.AUTOTUNE).batch(batch_size)
    dataset = dataset.prefetch(buffer_size=tf.data.experimental.AUTOTUNE)
    return dataset

现在可以运行到97%了：

一个epoch可以24S了：

内存不够，不用cache

这种情况会慢一点，GPU利用可以达到90%左右：

但是不占内存，一个epoch大概32S：

batch, shuffle, repeat这三者的顺序

batch, shuffle, repeat这三者的顺序不同，结果也会不一样。具体可以运行下面代码并调整三者顺序来验证

    tf.compat.v1.enable_eager_execution()
    dataset = tf.data.Dataset.from_tensor_slices(list(range(0, 10)))
    print('init ds:',list(dataset.as_numpy_iterator()))
    dataset = dataset.batch(3, drop_remainder=False)
    print('batched ds:',list(dataset.as_numpy_iterator()))
    dataset = dataset.repeat(2)
    print('batched + repeated ds:',list(dataset.as_numpy_iterator()))
    dataset = dataset.shuffle(4)
    print('batched + repeated + shuffle ds:',list(dataset.as_numpy_iterator()))

得到：

init ds: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
batched ds: [array([0, 1, 2]), array([3, 4, 5]), array([6, 7, 8]), array([9])]
batched + repeated ds: [array([0, 1, 2]), array([3, 4, 5]), array([6, 7, 8]), array([9]), array([0, 1, 2]), array([3, 4, 5]), array([6, 7, 8]), array([9])]
batched + repeated + shuffle ds: [array([6, 7, 8]), array([9]), array([0, 1, 2]), array([6, 7, 8]), array([0, 1, 2]), array([9]), array([3, 4, 5]), array([3, 4, 5])]