如何在TensorFlow中并行化for循环以并行训练多个模型？

我目前在TensorFlow中通过for循环训练多个独立模型，但训练过程是串行执行的，没办法充分利用GPU的并行计算能力。我的场景和Stack Overflow上一个未得到解答的问题类似，但我是在训练阶段（而非仅推理）遇到这个问题。那个问题下有一条关键评论：

提升速度的方法是使用单个进程与单个TensorFlow运行时。TensorFlow未针对多个运行时共享GPU的场景做优化

我的测试观察 #

我做了一组对比测试，结果差异很明显：

• 当我构建一个总参数量比循环中4个模型之和还大的单个模型时，每步训练耗时约0.19秒
• 而用for循环依次训练这4个不同模型时，整个循环每步耗时约0.55秒——几乎是单个大模型的4倍

所有模型都能完全放入GPU显存，我希望找到方法让for循环的所有迭代并行执行，使整个循环每步耗时接近单个大模型的0.19秒。

限制条件 #

由于工作性质，我无法使用tf.while_loop，因为它要求损失函数必须在while_loop内部定义。我想知道是否有其他方法在GPU上并行化这个for循环？比如使用TensorFlow的threading模块？另外我也没有设置per_process_gpu_memory_fraction，因为之前的问题指出这会降低训练速度。

额外疑问 #

另外，在Keras或PyTorch中能否实现这类多个独立模型的并行训练？

恳请各位提供帮助，谢谢！

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