I’d like to partition a neural network into two sub-networks using Pytorch. To make things concrete, consider this image:

partition_networks762×594 48.9 KB

In 1, I’ve a 3x4x1 neural network. What I want is, for example during epoch 1, I’d only like to update the weights in the sub-network 1, i.e., the weights that appear in the sub-network 2 must be frozen. Then again, in epoch 2, I’d like to train the weights that appear in sub-network 2 while the rest should frozen.
How can I do that?