TFRecord文件训练速度调优

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()方法。