1

我对在 TensorFlow 中使用 Dataset API 的性能指南有疑问。我有以下代码:

output_path =  "train_data_full_encoded_4.tfrecords"
batch_size  = 16
output_buffer_size = batch_size * 100
dataset            = tf.data.TFRecordDataset(output_path)
src_tgt_dataset    = dataset.shard(1, 0)
src_tgt_dataset    = src_tgt_dataset.shuffle(output_buffer_size, 44)
src_tgt_dataset    = src_tgt_dataset.map(_parse_record, num_parallel_calls=4).prefetch(output_buffer_size)
src_tgt_dataset    = src_tgt_dataset.map(_random_crop_filp, num_parallel_calls=4).prefetch(output_buffer_size)
src_tgt_dataset    = src_tgt_dataset.map(_make_labels_small, num_parallel_calls=None).prefetch(output_buffer_size)

import time
dataset = src_tgt_dataset.batch(batch_size)
dataset = dataset.repeat(1)
end    = time.time()
with tf.Session(config=tf.ConfigProto(allow_soft_placement=True)) as sess:
    batched_iter = dataset.make_initializable_iterator()
    images, scores, geomaps, trainingmask = batched_iter.get_next()
    sess.run(batched_iter.initializer)
    idx = 0
    while True:
        try:
            raw_images, raw_scores, raw_geomaps, raw_trainingmask = sess.run([images, scores, geomaps, trainingmask])

        except tf.errors.OutOfRangeError:
            break
print  idx, time.time() - end

我的问题是读取数据的速度很慢,每批大约需要 0.2 秒,这使得 GPU 的利用率很低。我想加快速度,这就是我创建 TfRecords 的原因。它比普通加载器运行得更快,但我仍然无法始终获得 >60% 的利用率。

如何加快数据加载速度?也许有一些性能指南?我很好奇设置num_parallel_calls不同的map功能是否对最终速度有任何影响(我尝试了几个版本,当前版本似乎是最快的)。

此外,为了创建 TfRecord,我使用编码为 JPEG 和 PNG 的图像,但没有使用tf.image.encode_jpegtf.image.encode_png因为我在使用它们时存在内存问题(这个过程占用了所有 16GB,但是当PIL 用于编码图像时,内存使用是正确的),有什么提示吗?这是我使用 TF 创建 TfRecords 的代码。

def createRecord(folder_with_annotation, idx, files):
    output_path = os.path.join(folder_with_annotation, "train_data_full_encoded_{}.tfrecords".format(idx))
    writer      = tf.python_io.TFRecordWriter(output_path)
    with tf.Session() as sess:
        for f in tqdm(files):
            img, score_map, geo_map, training_mask = preprocessImage(f)
            h, w, c = img.shape

        image_tf      = tf.image.encode_jpeg(tf.constant(img),format='rgb').eval(session=sess)

        score_map_tf     = tf.image.encode_jpeg(tf.constant(score_map),format='grayscale').eval(session=sess)
        training_mask_tf = tf.image.encode_jpeg(tf.constant(training_mask),format='grayscale').eval(session=sess)
        geo_map_part_one = tf.image.encode_png(tf.constant(geo_map)).eval(session=sess)

        example = tf.train.Example(features=tf.train.Features(feature={
                    'height': _int64_feature(h),
                    'width': _int64_feature(w),
                    'depth': _int64_feature(c),
                    'multiply_geo_map': _int64_feature(multiply),
                    'image_raw': _bytes_feature(image_tf),
                    'score_map_raw': _bytes_feature(score_map_tf),
                    'geo_map_raw': _bytes_feature(geo_map_part_one),
                    'training_mask_raw': _bytes_feature(training_mask_tf),
                    }))

        writer.write(example.SerializeToString())

# All examples saved
writer.close()

编辑:我在创建 TfRecord 时解决了内存问题,但不会增加图形(并运行eval完成)。tf.encode使用功能时,我会在得到阅读速度的结果后通知。Edit2:正如我所料,这不会减少读取数据的时间(但 TfRecord 的大小更小)。

4

0 回答 0