A model with multiple outputs

Multiple outputs is pretty straightforward. Just return mutiple values in the forward() method of your net.

def forward(self, x):
    # Do your stuff here
    ...
    x1 = F.log_softmax(x) # class probabilities
    x2 = ... # bounding box calculation
    return x1, x2

Using these two outputs, you can define two different loss functions and just add them.

out1, out2 = model(data)
loss1 = criterion1(out1, target1)
loss2 = criterion2(out2, target2)
loss = loss1 + loss2
loss.backward()
46 Likes