Fine-Tune DETR on custom dataset (less than 250 labels)

Models
1

I’ve been following the wonderful tutorials by @NielsRogge to fine-tune DETR on a custom dataset. However, my predicted mask seems to be always empty after training.

Let’s say my dataset has 4 labels (in coco panoptic format). Looking at the model.config.id2label of the pretrained DETR there are 250 different labels.

What would be the correct way to fine-tune the model to less labels (similar to fine-tuning a ViT classifier to 2 classes)?

Here is what I tried:

num_labels = 4

feature_extractor = DetrFeatureExtractor.from_pretrained(
    "facebook/detr-resnet-50-panoptic", size=224, format="coco_panoptic"
)

config = DetrConfig.from_pretrained(
    "facebook/detr-resnet-50-panoptic",
    num_labels=num_labels
)
model = DetrForSegmentation.from_pretrained("facebook/detr-resnet-50-panoptic")
model.detr.class_labels_classifier = Linear(
    in_features=model.config.hidden_size,
    out_features=num_labels + 1,  # +1 for "no object" class
)

# remove classifier weights and bias for pretrained model, since it has more classes than we want
state_dict = model.state_dict()
state_dict.pop("detr.class_labels_classifier.weight")
state_dict.pop("detr.class_labels_classifier.bias")
model.load_state_dict(state_dict, strict=False)

model.config = config

Even though my loss goes down, it appears that the model doesn’t learn. The post-processed predictions using feature_extractor.post_process_panoptic() show the following (empty) result:

{'png_string': b'\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x01\xdf\x00\x00\x00\xe0\x08\x02\x00\x00\x00\xcd9\xfb\xc3\x00\x00\x01OIDATx\x9c\xed\xc11\x01\x00\x00\x00\xc2\xa0\xf5Om\x0c\x1f\xa0\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\xbe\x06\xea|\x00\x01\xd1K+\xbf\x00\x00\x00\x00IEND\xaeB`\x82',
 'segments_info': []}

For feature_extractor.post_process_segmentation() I get the following:

{'scores': tensor([], grad_fn=<IndexBackward0>),
 'labels': tensor([], dtype=torch.int64),
 'masks': tensor([], size=(0, 426, 640), dtype=torch.int64)}

Side Question: Would it make sense to lower the number of queries (from default 100) to something like 10 (from DetrFeatureExtractor documentation: “Note that it’s good to have some slack”)?

2

The default threshold in post_process_segmentation() is .50. Try reducing that ?
On a separate note any insights on training with a small dataset ?