我对在 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_jpeg
,tf.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 的大小更小)。