Custom dataset training

[facebookresearch]

Custom dataset training

[fmassa]

Hi,

The error you are facing is because the number of classes is required for creating the criterion, see

detr/models/detr.py

Line 90 in b7b62c0

self.num_classes = num_classes

and

detr/models/detr.py

Lines 95 to 97 in b7b62c0

	empty_weight = torch.ones(self.num_classes + 1)
	empty_weight[-1] = self.eos_coef
	self.register_buffer('empty_weight', empty_weight)

you can probably just overwrite those values with the number of classes you want, and this error should disappear.

Hi,

Thank you! It worked. Now I am training the class_embed and bbox_embed layers for an additional class 'skyscraper'. The goal is to add additional classes to the 91 classes in the pre-trained model. Hopefully, it will work.

Thanks!

[fmassa]

Great, keep us informed on how it works.

I believe I have answered your question, and as such I'm closing the issue, but please let us know if you have further questions / issues