我试图优化我的数据输入管道.
以下是管道:
def build_dataset(file_pattern):
return tf.data.Dataset.list_files(
file_pattern
).interleave(
tf.data.TFRecordDataset,
num_parallel_calls=tf.data.experimental.AUTOTUNE
).shuffle(
buffer_size=2048
).batch(
batch_size=2048,
drop_remainder=True,
).cache(
).repeat(
).map(
map_func=_parse_example_batch,
num_parallel_calls=tf.data.experimental.AUTOTUNE
).prefetch(
buffer_size=1
)
使用映射函数:
def _bit_to_float(string_batch: tf.Tensor):
return tf.reshape(tf.math.floormod(tf.dtypes.cast(tf.bitwise.right_shift(
tf.expand_dims(tf.io.decode_raw(string_batch, tf.uint8), 2),
tf.reshape(tf.dtypes.cast(tf.range(7, -1, -1), tf.uint8), (1, 1, 8))
), tf.float32), 2), (tf.shape(string_batch)[0], -1))
def _parse_example_batch(example_batch):
preprocessed_sample_columns = {
"features": tf.io.VarLenFeature(tf.float32),
"booleanFeatures": tf.io.FixedLenFeature((), tf.string, ""),
"label": tf.io.FixedLenFeature((), tf.float32, -1)
}
samples = tf.io.parse_example(example_batch, preprocessed_sample_columns)
dense_float = tf.sparse.to_dense(samples["features"])
bits_to_float = _bit_to_float(samples["booleanFeatures"])
return (
tf.concat([dense_float, bits_to_float], 1),
tf.reshape(samples["label"], (-1, 1))
)
我试图遵循data pipeline tutorial的最佳实践,并将我的映射函数矢量化(正如mrry建议的那样).
在这种设置下,当数据以高速(带宽约为200MB/s)下载时,CPU使用不足(14%),训练非常缓慢(一个历元超过1小时).
我try 了一些参数配置,改变了interleave()
个参数,比如num_parallel_calls
或cycle_length
,或者TFRecordDataset
个参数,比如num_parallel_calls
.
最快的配置使用这组参数:
-
interleave.num_parallel_calls
: 1 -
interleave.cycle_length
: 8 -
TFRecordDataset.num_parallel_calls
: 8
有了这个,一个历元只需要20分钟就可以运行.However, CPU usage is only at 50% while bandwidth consumption is around 55MB/s
问题:
- 如何优化流水线以达到100%的CPU使用率(以及大约100MB/s的带宽消耗)?
- 为什么
tf.data.experimental.AUTOTUNE
没有找到加速训练的最佳价值?
友善的
编辑
经过更多的实验,我得出了以下解决方案.
- 如果
num_parallel_calls
大于0,则删除TFRecordDataset
已经处理的interleave
步. - 将映射函数更新为仅执行
parse_example
和decode_raw
,返回元组"(,),()) -
map
后cache
- 将
_bit_to_float
功能作为模型的一个组件移动
最后,这里是数据管道代码:
def build_dataset(file_pattern):
return tf.data.TFRecordDataset(
tf.data.Dataset.list_files(file_pattern),
num_parallel_reads=multiprocessing.cpu_count(),
buffer_size=70*1000*1000
).shuffle(
buffer_size=2048
).map(
map_func=split,
num_parallel_calls=tf.data.experimental.AUTOTUNE
).batch(
batch_size=2048,
drop_remainder=True,
).cache(
).repeat(
).prefetch(
buffer_size=32
)
def split(example):
preprocessed_sample_columns = {
"features": tf.io.VarLenFeature(tf.float32),
"booleanFeatures": tf.io.FixedLenFeature((), tf.string, ""),
"label": tf.io.FixedLenFeature((), tf.float32, -1)
}
samples = tf.io.parse_single_example(example, preprocessed_sample_columns)
dense_float = tf.sparse.to_dense(samples["features"])
bits_to_float = tf.io.decode_raw(samples["booleanFeatures"], tf.uint8)
return (
(dense_float, bits_to_float),
tf.reshape(samples["label"], (1,))
)
def build_model(input_shape):
feature = keras.Input(shape=(N,))
bool_feature = keras.Input(shape=(M,), dtype="uint8")
one_hot = dataset._bit_to_float(bool_feature)
dense_input = tf.reshape(
keras.backend.concatenate([feature, one_hot], 1),
input_shape)
output = actual_model(dense_input)
model = keras.Model([feature, bool_feature], output)
return model
def _bit_to_float(string_batch: tf.Tensor):
return tf.dtypes.cast(tf.reshape(
tf.bitwise.bitwise_and(
tf.bitwise.right_shift(
tf.expand_dims(string_batch, 2),
tf.reshape(
tf.dtypes.cast(tf.range(7, -1, -1), tf.uint8),
(1, 1, 8)
),
),
tf.constant(0x01, dtype=tf.uint8)
),
(tf.shape(string_batch)[0], -1)
), tf.float32)
感谢所有这些优化:
- 带宽消耗约为90MB/s
- CPU使用率约为20%
- 第一纪元花了20分钟
- Successives历险记每次花费5分钟
所以这似乎是一个很好的第一次设置.但是CPU和BW仍然没有被过度使用,所以任何建议都是受欢迎的!
编辑 Bis
所以,在进行了一些基准测试之后,我了解了我认为最好的输入管道:
def build_dataset(file_pattern):
tf.data.Dataset.list_files(
file_pattern
).interleave(
TFRecordDataset,
cycle_length=tf.data.experimental.AUTOTUNE,
num_parallel_calls=tf.data.experimental.AUTOTUNE
).shuffle(
2048
).batch(
batch_size=64,
drop_remainder=True,
).map(
map_func=parse_examples_batch,
num_parallel_calls=tf.data.experimental.AUTOTUNE
).cache(
).prefetch(
tf.data.experimental.AUTOTUNE
)
def parse_examples_batch(examples):
preprocessed_sample_columns = {
"features": tf.io.FixedLenSequenceFeature((), tf.float32, allow_missing=True),
"booleanFeatures": tf.io.FixedLenFeature((), tf.string, ""),
"label": tf.io.FixedLenFeature((), tf.float32, -1)
}
samples = tf.io.parse_example(examples, preprocessed_sample_columns)
bits_to_float = tf.io.decode_raw(samples["booleanFeatures"], tf.uint8)
return (
(samples['features'], bits_to_float),
tf.expand_dims(samples["label"], 1)
)
那么,最新消息是:
- 根据这个GitHub issue,
TFRecordDataset
交织是传统的,所以interleave
功能更好. -
map
之前的batch
是一个好习惯(vectorizing your function),并减少调用映射函数的次数. - 不再需要
repeat
了.自TF2以来.0时,Keras模型API支持dataset API并可以使用缓存(请参阅SO post) - 从
VarLenFeature
切换到FixedLenSequenceFeature
,取消对tf.sparse.to_dense
的无用呼叫.
希望这能有所帮助.建议仍受欢迎.