移植PyTorch版Gridworld DRL代码到TensorFlow后训练耗时过长，原因何在？

你遇到的性能差距完全来自实现层面的非向量化操作和TF运行时的配置问题，和算法逻辑本身无关，具体原因如下：

• 头号元凶：训练循环内的Python级逐元素索引
  对比两版代码中动作值提取的逻辑：
  PyTorch用Q1.gather()实现了全C层的向量化批量动作选择，没有任何Python循环开销。
  你的TensorFlow版用列表推导[Q1[i][action_batch[i]] for i in range(len(action_batch))]做逐元素索引，每一次张量索引都要触发Python和TF C运行时的交互，在CPU上这种Python循环比向量化张量操作慢1~2个数量级，是造成几小时训练时长的核心原因。
• 未使用TF原生的向量化索引接口
  TF中对应PyTorch gather的批量动作选择不需要任何Python循环，直接使用tf.gather配合batch_dims=1参数即可实现完全等价的向量化计算：

  # 替换原来的列表推导代码
  X = tf.gather(Q1, action_batch, axis=1, batch_dims=1)
  

  替换后这部分的性能和PyTorch基本持平，计算结果完全一致。
• Eager模式的固有调度开销
  TensorFlow默认的eager执行模式下，每一步张量运算都要经过Python层调度，即使没有循环，整体运算调度开销也比PyTorch略高。可以通过tf.function把单步训练逻辑编译成静态图，进一步降低CPU上的调度开销，通常能提速30%~50%。
• 隐性的线程配置开销
  部分环境下TF默认的CPU线程数配置不合理，会导致无意义的线程上下文切换。可以在脚本开头添加如下配置，指定匹配物理核心数的并行线程数：

  # 替换为你机器的物理CPU核心数
  tf.config.threading.set_intra_op_parallelism_threads(4)
  tf.config.threading.set_inter_op_parallelism_threads(2)
  

以上修改完成后，你的TF版训练速度基本可以和PyTorch版持平，最终胜率不会出现明显波动。

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