Using a different box predictor with FasterRCNN

I am trying to solve an object detection problem using FasterRCNN, where I have N detection scenarios. I am getting good results with transfer learning just the last layer (box predictor), but I end up having N models, one for each scenario. Consequently, it takes up a lot of storage space and RAM if I want all the models to be loaded simultaneously. Now, only the box predictor (last layer) is different for each network, so I want to pass my inputs in a batch, get the output from forward pass till only the last layer, and then pass the individual outputs to N different box predictors. What’s the best strategy to implement this? I can only think of swapping the last layer before running on individual inputs, which is too slow.


Hi @ssetu,

I think the best way to achieve that would be to create a Module inheriting from FasterRCNN (minus the box_predictor) with a list of your fine-tuned box_predictors, define a custom forward with an extra “detection scenario” argument where you would switch on the right box_predictor, depending on that extra argument.

1 Like

Thank you @spanev. I will try this out with clues from Add multiple FC layers in parallel and post an update soon.

1 Like