How to train a custom model for object detection using models/official/vision/detection?

Viewed 233
1 Answers

To train a new model, the training entry is main.py.

Here are a few steps of how to add new models.

If you want to just build a simple model, say MyRetinaNet, on top of current existing components like layers, losses, existing heads, you might need to:

  1. Add a config template for the new model like this one.
  2. Add a file "my_retinanet_model.py" in the modeling folder (similar to "retinanet_model.py") and implement the model.
  3. Add a branch to the factory file so that you can use it in main.py.

If you want to add some fine-grained components like heads and backbones, then you need to add something to the models/official/vision/detection/modeling/architecture/ folder.

  1. Add a class to heads.py (for heads) or a new .py file for backbones.
  2. Update the factory.py accordingly.
  3. You might also need to change the model configs accordingly,

For even finer-grained ops, you can add to ops and utils.

Related