How to Add a Regression Head

Question

I am fairly comfortable using tensorflow directly (python) as well as building models with keras. So far all the models I have built have been traditional straight through processing from the input image, through to the softmax output. I'm actually surprised how straightforward it is.

I now want to train a model for both classification and regression. Slide 16 of this Stanford CNN class CS231n is exactly what I want to do. I'm really unclear (actually, I have NFI) how to branch the tensorflow graph, create two loss functions and properly back propagate both of them.

Can anyone point me to a code sample? Very much appreciated.

Answer 1

Responding to the question in the comment:

I have not worked on such problems but here are some general thoughts.

Ideally, the (shared) weights that are good for regression should also be good for classification. In this case, the loss_weights don't matter much as both losses should push the weights in the same direction. You should just make sure that the average magnitudes of your loss * loss_weight is roughly the same so that one loss is not completely ignored.

If you find that the heads are fighting, ie you can train a separate regression model to much higher accuracy than if it is trained together with classification, you can either branch-out the heads earlier (so that they share fewer weights and have more specialized layers for themselves) or just train two separate models. There are also a ton of research papers on multi-task training that allow you to share more weights to accomplish different tasks.