How to get QueryInst to train on an external dataset

[fbragman]

Hi,

I am trying to train QueryInst with a Swin backbone on my own medical imaging temporal dataset which contains 3 classes - background + 2 anatomical landmarks. The reason being I would then like to initialise a TeViT model with Swin-QueryInst weights to help the temporal model.

To initialise the Swin backbone in the QueryInst training procedure, I pretrained a segmentation model using Swin + some output segmentation layers then initialise the backbone with the pretrained weights. I set the learning rate of the backbone to be 0.1 * lr of the ROI head + Bbox heads. Additionally - I use a StepLR schedulers and similar AdamW parameters to those published in the paper. Moreover, I used gradient clipping with similar values (norm=1, type=2)

However, performance of the QueryInst model is really poor and the segmentation performance as measured through IoU per class degrades significantly from the baseline I trained.

The mAP and mAP_0.5 in the training set seems to converge nicely. However, the models fails to learns a robust function for the instance masks.

Do you have any suggestions?