In pytorch, how to train a model with two or more outputs?

Question

output_1, output_2 = model(x)
loss = cross_entropy_loss(output_1, target_1)
loss.backward()
optimizer.step()

loss = cross_entropy_loss(output_2, target_2)
loss.backward()
optimizer.step()

However, when I run this piece of code, I got this error:

RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation: [torch.FloatTensor [1, 4]], which is output 0 of TBackward, is at version 2; expected version 1 instead. Hint: enable anomaly detection to find the operation that failed to compute its gradient, with torch.autograd.set_detect_anomaly(True).

Then, I really wanna know what I am supposed to do to train a model with 2 or more outputs

Answer 1

The entire premise on which pytorch (and other DL frameworks) is founded on is the backporpagation of the gradients of a scalar loss function.
In your case, you have a vector (of dim=2) loss function:

[cross_entropy_loss(output_1, target_1), cross_entropy_loss(output_2, target_2)]

You need to decide how to combine these two losses into a single scalar loss.
For instance:

weight = 0.5  # relative weight
loss = weight * cross_entropy_loss(output_1, target_1) + (1. - weight) * cross_entropy_loss(output_2, target_2)
# now loss is a scalar
loss.backward()
optimizer.step()