TFRecord文件训练速度调优
-
问题描述：
-
背景：
-
tfrecord文件由spark生成：
res_df.write.format("tfrecords").option("recordType","Example")mode("overwrite").save("path/") -
大约有90个特征，其中有int型、string型、单值序列型、多值序列型（保存时拍平为一维列表了，读取时指定二维维度
-
解析时使用tfrecorddataset进行解析
-
训练时将使用fit()接口，将tfrecorddataset传入
-
tensorflow版本：2.6.0
-
-
问题：
- 速度缓慢
- GPU利用率低
-
-
调优关键代码：
ds = tf.data.TFRecordDataset(p) # p为tfrecord文件路径，因为只有一个文件，不要设置num_parallel_reads ds = ds.shuffle(ds_batch_size) # shuffle_size设置为batch_size ds = ds.batch(ds_batch_size, num_parallel_calls=64) # 先batch，再map解析函数，num_parallel_calls设置为核心数，不是线程数 ds = ds.map(m_decode_fn, num_parallel_calls=64) # 一样，设置为核心数，超过会争夺资源，引发阻塞 # ds = ds.prefetch(1) ds = ds.apply(tf.data.experimental.prefetch_to_device('/gpu:0', buffer_size=4)) # prefetch改为直接送到GPU，可以减少大量传输时间 -
其他可关注点：
- 使用profile+tensorboard可以看到整个数据流各个环节的时间耗时，并且可以得到部分优化建议。
- 可以将单个tfrecord文件分为多个tfrecord文件，然后使用num_parallel_reads>1的方法去多线程地读。我这里没啥区别，主要耗时还是在解析和cpu-gpu数据复制上。
- 数据集比较小（至少内存装得下），可以试试cache()方法。