为何PyTorch的for循环比TensorFlow2.3快10倍？附代码求优化方案

• eager执行开销：TensorFlow 2.3的eager执行模式在Python侧循环中运行时，每一步Tensor操作都需要在Python解释器和TF内核间切换，调度开销远高于PyTorch优化后的eager执行。你的场景是自回归序列生成，循环每一步依赖上一步输出，无法被TF自动向量化，开销会随seq_len累加。
• RNN调用冗余：TensorFlow代码中每次都对单时间步输入做tf.expand_dims(x, axis=1)后传入完整RNN层，RNN层会额外执行时间维度处理、状态校验等冗余逻辑；而PyTorch代码直接调用RNN的单步状态更新逻辑，没有额外开销。
• 设备同步开销：循环中i % 100 == 0分支使用了Python原生的time.time和自定义回调，会强制触发CPU和GPU的设备同步，每次同步都需要等待GPU上所有操作执行完成，额外增加了耗时。
• 采样逻辑性能差：TensorFlow Probability的Categorical分布在2.3版本的eager模式下初始化和采样性能远低于PyTorch的对应实现，每次循环新建分布对象的开销不可忽视。

• 用静态图编译整个循环逻辑：将包含循环的函数用@tf.function装饰，编译为静态计算图，消除Python侧的调度开销。注意循环中存储输出的Python list需要替换为tf.TensorArray，静态图不支持原生list的append操作。
• 替换为RNN Cell单步调用：不要使用完整RNN层处理单时间步输入，直接调用RNN层的cell属性执行单步更新，写法对齐PyTorch版本：_, h1 = rnn1.cell(x, h1)，省去冗余的时间维度处理开销。
• 简化采样逻辑：不要每次新建Categorical分布对象，直接使用tf.random.categorical(logits, num_samples=1)完成采样，大幅降低采样开销。
• 移除/改造调试逻辑：如果不需要实时进度反馈，直接删除i % 100 == 0的分支逻辑；如果需要保留，将time.time替换为tf.timestamp，回调函数用@tf.autograph.experimental.do_not_convert装饰，避免频繁触发设备同步。

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