How to keep the outputs of different networks and then use it?

#1

I need to aggregate the outputs from different networks and then do some calculations based on this aggregated outputs, and finally update individual networks. I am unclear about how to keep the outputs which are obtained in the first for loop (from individual networks) and use it in the second for loop to calculate the loss and update individual networks. I very much appreciate if you give me any suggestion.

outputs = []
for i in range(noOfModels):
    optimizer[i].zero_grad()
    tmp = net[i](img) # this will return an output of size 10x10 
    outputs.append(tmp)

outputs = torch.stack(outputs, 2)
x = do_somthing(outputs) # do some calculations with outputs

for i in range(noOfModels):
    loss = criterion(outputs[:,:,i], x)
    loss.backward()
    optimizer[i].step()
#2

I would probably just add the losses like this:

loss = 0
for i in range(noOfModels):
    loss = loss + criterion(outputs[:,:,i], x)
loss.backward()

for i in range(noOfModels):
    optimizer[i].step()

This should work if you didn’t break the computational graph somewhere.

1 Like