facebookresearch / detr

End-to-End Object Detection with Transformers
Apache License 2.0
13.55k stars 2.45k forks source link

support custom dataset training with args.num_classes #88

Open lessw2020 opened 4 years ago

lessw2020 commented 4 years ago

❓ How to do something using DETR

Describe what you want to do, including:

  1. what inputs you will provide, if any:
  2. what outputs you are expecting:

Users want to be able to easily train DETR via fine tuning for their own datasets. Each user dataset will have it's own num_classes.

Currently the detr.py::build(args) function determines the num_classes by defaulting to 20 or 91 (coco) or 250 (coco_panoptic), based on the args.dataset_file param:

https://github.com/facebookresearch/detr/blob/master/models/detr.py#L304

Thus, request: is there an addition to support an args.num_classes to be passed in that can be easily adjusted, and have that be used for model creation instead of the defaults based on dataset_file arg?

I could make a quick PR for this if that helps?

NOTE:

  1. Only general answers are provided. If you want to ask about "why X did not work", please use the Unexpected behaviors issue template.

  2. About how to implement new models / new dataloader / new training logic, etc., check documentation first.

  3. We do not answer general machine learning / computer vision questions that are not specific to DETR, such as how a model works, how to improve your training/make it converge, or what algorithm/methods can be used to achieve X.

lessw2020 commented 4 years ago

update - I made a PR to hopefully cleanly integrate this support with full backwards compat: https://github.com/facebookresearch/detr/pull/89