Recommendations for training Detr on custom dataset?

[lessw2020]

Very impressed with the all new innovative architecture in Detr!
Can you clarify recommendations for training on a custom dataset?
Should we build a model similar to demo and train, or better to use and fine tune a full coco pretrained model and adjust the linear layer to desired class count?
Thanks in advance for any input.

[Zumbalamambo]

+1

[PancakeAwesome]

agree

[alcinos]

Hello,
Thanks for your interest in DETR.
It depends on the size of your dataset. If you have enough data (say at least 10K), training from scratch should work just fine. You'll need to prepare the data in the coco format and then follow instructions from the Readme. Note that if your dataset has a substantially different average number of objects per image than coco, you might need to adjust the number of object queries (--num_queries) It should be strictly higher than the max number of objects you may have to detect, and it's good to have some slack (in coco we use 100, the max number of objects in a coco image is ~70)

Fine-tuning should work in theory, but at the moment it's not tested/supported. If you want to give it a go anyways, you just need to --resume from one of the checkpoint we provide. Feel free to report back any results you obtain :)

Best of luck

[raviv]

Hi,

When fine-tuning from model zoo, using my own dataset, how should I modify the number of classes?
Loading the model fails (as expected) on:

RuntimeError: Error(s) in loading state_dict for DETR:
	size mismatch for class_embed.weight: copying a param with shape torch.Size([92, 256]) from checkpoint, the shape in current model is torch.Size([51, 256]).
	size mismatch for class_embed.bias: copying a param with shape torch.Size([92]) from checkpoint, the shape in current model is torch.Size([51]).

As I have 50 labels, and the checkpointed model has 91.

Thanks!

[alcinos]

If you just want to replace the classification head, you need to erase it before loading the state dict. One approach would be:

model = torch.hub.load('facebookresearch/detr', 'detr_resnet50', pretrained=False, num_classes=50)
checkpoint = torch.hub.load_state_dict_from_url(
            url='https://dl.fbaipublicfiles.com/detr/detr-r50-e632da11.pth',
            map_location='cpu',
            check_hash=True)
del checkpoint["model"]["class_embed.weight"]
del checkpoint["model"]["class_embed.bias"]
model.load_state_dict(checkpoint["model"], strict=False)

Best of luck.

[cbasavaraj]

It would be easier (or at least more standard practice) to first load the pre-trained model, and then replace the classification head.

[lessw2020]

related question but how should we downgrade the query number for smaller classes ( in terms of continuing from the approach above)?
For example I only have 5 classes to detect and each image will have exactly 5 classes per image, so I was planning to run with queries = 12 instead of the default 100 (or should it be 5 if we know that's the max our images will ever have...)

I'm looking at model.query_embed with (100,256) and assume that is the right place to adjust but unclear. If we adjust via model.query_embed.num_embeddings=my_new_query_count, is that enough?
(update - I'm working on this and the DETR model stores a self.num_queries as well, but this is only referenced later for segmentation.
But to be correct should update both model.num_queries and the model.query_embed.num_embeddings would need to be adjusted together...)

[lessw2020]

Also wouldn't we want to re-init the weights in class_embed to normal or uniform after wiping the checkpoint weights to kick off the new training?

[alcinos]

If you're fine-tuning, I don't recommend changing the number of queries on the fly, it is extremely unlikely to work out of the box. In this case you're probably better off retraining from scratch (you can change the --num_queries arg from our training script).

As for the initialization of class_embed, the solution I posted above makes sure it is initialized as it should.

Best of luck

[lessw2020]

Hi @alcinos - excellent, thanks tremendously for the advice here, esp on a Sat night.
I will try both fine tuning for now (with smaller dataset and will not touch num_queries) and from scratch as we'll have a larger dataset soon, and update here to share results.
Thanks again!

[raviv]

My dataset has images of various sizes.
Do I need to resize them to a specific size?