Assume we have two pytorch models M1 and M2. Each one are made of two fully connected layers. The first layer in these two models are shared.
The models are trained in parallel. A batch of data is fed into the first layer and then the output is fed into the second layer of each network to produce o1 and o2 (i.e., the outputs of the first and second networks).
We define the loss function L to be:
L1 = Criterion(o1, y1)
L2 = Criterion(o2, y2)
L = L1 + L2 + Diff_grad
y1 and y2 are the true labels for the two networks and Criterion() is binary cross entropy.
Diff_grad is the difference between the gradients of the two losses with respect to the first layer in the networks. More specifically it is: (gradient_of_L1_wrt_layer1 - gradient_of_L2_wrt_layer1)^2.
In plain language, I am trying to minimize the two losses, as well as the difference between the gradients of the two losses with respect to the first layer.
Can anyone tell me how to do this (preferably by a piece of code)? Thanks.
PS. the actually task that I am doing is bigger than this, the networks are much bigger, and they have more shared modules. Also the objective is not the “difference between the gradients”, it is a more complicated function. The example above is for my understanding.