Refactor the `forward` method of the RoiHeads into separate methods for easier customization

🚀 Feature

Break up the code in the forward method into several sub-methods and allow custom loss and inference functions.

Motivation

I have a use case that requires implementing a different loss function for masks during training a Mask-RCNN, but designing this is somewhat messy right now. Ideally, I'd like to be able to inherit RoiHeads and then make a small change to a method like process_masks rather than overriding the whole forward method. (Or even simpler, just pass in custom loss functions during initialization).

Pitch

The structure of forward lends itself to re-factoring really simply.
forward could call, in sequence:
process_boxes(...)
process_masks(...)
process_keypoints(...)

This way, changing the behaviour requires only re-implementing one of these methods in the child class rather than all of forward. Optionally, these functions could use customized loss and inference functions, maybe passed into RoiHeads during initialization.

Alternatives

For my use case, it would also be sufficient to be able to optionally pass RoiHeads custom loss functions during initialization, but I think breaking it up is still a good idea because it makes other changes that do require changing how some of these processes are handled a little simpler and cleaner.

Another option would be to add the functions related to loss and inference to the class itself, so they could be directly overriden by a child class, but I think this may muddy the responsibilities of RoiHeads a little.

Additional context

Direct link to module: https://github.com/pytorch/vision/blob/master/torchvision/models/detection/roi_heads.py
I'm happy to take this on if approved.

Start with torchvision/models/detection/roi_heads.py and read the RoiHeads forward method and surrounding box, mask, and keypoint processing. Determine how the proposed sub-methods and customizable loss or inference functions fit the existing API, then verify that the refactoring preserves current training and inference behavior.