Hi everyone.
I have a model that is built on the top of multiple nn.Modules and my model have two loss functions. The error of one of these two loss functions will have to backpropagate through the whole network while the second error should only update one of the modules that build the model.
As an example, my modules are:
class first_module(nn.Module):
#some code here#
class second_module(nn.Module):
#some code here#
My final network is something like:
class model(nn.Module):
def __ init __(self):
self.first_part = first_module()
self.second_part = second_module()
def forward(input):
x = self.first_part (input)
y = self.second_part(x)
z = some_operations(y)
return z,y
Now on my main code i have:
network = model()
output_1, output_2 = network(image)
loss_1 = error_fn(output_1 , labels) # main loss
loss_2 = error_fn(output_2, ref) # i want this error to only update layers on (class second_module)
In other words, second_module is updated jointly by loss_1 and loss_2, but other parts of the modules are only updated by loss_1.
Your help is highly appreciated.