parallel-processing - 如何并行拆分和训练不同的模型部分（pytorch）？

Question

我有这个模型结构： 在此处输入图像描述

我在拆分模型之前将所有层都写到同一个类中，但是我想训练如图所示，我拆分了类但发现反向发送时只更新了最后一个过程中的参数4，结果不好。有什么方法可以实现我的训练模式？我认为它与 torch.multiprocessing 或 DP 或 DDP 不同。