TensorFlow2模型训练过慢 如何检测网络各组件实际运行耗时

直接上可落地的方法，避开异步执行导致的统计偏差：

• 计算图内嵌原生时间戳打点（精度最高，适合单组件定向统计）
  不要用Python原生datetime/time.time()打点，这类接口只能统计Python主线程的调度时间，不会等待GPU/TPU等异步设备的计算完成，结果完全不准。直接在计算图中插入TF原生时间戳算子，配合控制依赖保证执行顺序，就能拿到设备侧的真实执行时间。
  可以封装成无侵入的统计层，直接插在需要统计的模块前后：

  import tensorflow as tf
  from tensorflow.keras.layers import Layer
  import time
  
  class TimeMarker(Layer):
      def __init__(self, tag):
          super().__init__(name=f"time_marker_{tag}")
          self.tag = tag
      def call(self, inputs):
          # 插入设备侧时间戳算子
          ts = tf.timestamp()
          # 把时间戳作为metric上报，不影响原计算流
          self.add_metric(ts, name=f"ts_{self.tag}")
          return inputs
  
  # 用法示例：统计注意力模块耗时
  x = TimeMarker("attn_start")(input_tensor)
  x = MyAttentionBlock()(x)
  x = TimeMarker("attn_end")(x)
  

  两个marker的时间差就是对应模块的真实执行耗时，统计前先跑3-5步warmup，跳过图编译、显存分配、算子autotune的一次性开销。

• 正确配置tf.profiler追踪设备侧算子耗时
  之前用tf-profiler只看到Python层耗时，是因为默认配置没开设备内核追踪，按如下参数启动即可拿到完整的算子级时间线：

  # 先warmup，排除初始化开销
  for _ in range(5):
      model.train_on_batch(next(train_dataset))
  
  # 启动带设备追踪的profiler，只统计10-20步即可，避免日志过大
  tf.profiler.experimental.start(
      logdir="./prof_log",
      options=tf.profiler.experimental.ProfilerOptions(
          host_tracer_level=2,
          python_tracer_level=1,
          device_tracer_level=1, # 核心参数：开启GPU/CPU设备侧内核执行追踪
      )
  )
  for _ in range(15):
      model.train_on_batch(next(train_dataset))
  tf.profiler.experimental.stop()
  

  生成的日志可以直接看到每个算子在设备上的排队、执行、等待数据的精确耗时，能直接定位是算子本身计算慢、数据IO阻塞还是跨设备同步拖慢速度。

• 全局同步屏障快速统计相对耗时
  如果不想改动模型结构，可以在打点前调用tf.test.experimental.sync_devices()强制所有设备的计算任务执行完成，再用Python时间接口打点，这种方法会引入额外的同步开销，统计的绝对耗时会略高于实际训练值，但用来对比各模块的耗时占比、快速定位瓶颈完全够用：

  # 单步统计示例
  tf.test.experimental.sync_devices()
  t0 = time.time()
  # 执行要统计的模块/训练步
  with tf.GradientTape() as tape:
      loss = model(x, training=True)
  grads = tape.gradient(loss, model.trainable_variables)
  opt.apply_gradients(zip(grads, model.trainable_variables))
  tf.test.experimental.sync_devices()
  print(f"step cost: {time.time()-t0:.4f}s")
  

注意：所有统计都必须排除首次执行的开销，tf.function第一次运行时会做图追踪、编译、算子选型，耗时是稳态运行的数倍到数十倍，不属于正常训练的耗时范围。

内容的提问来源于stack exchange，提问作者zuijiang