Classes of the trained models?

[pieris98]

Hey Nazir,
Quick questions:

1. Are all the trained models in the RbA Model Zoo using (trained on) only the 19 STUFF classes (semantic segmentation) defined in the mask2former/datasets/data/register_cityscapes_partial.py file?
2. If that is the case, let's say I want to re-train/fine-tune RbA with some additional cityscapes classes (e.g. rail track). Is possible to re-train with the provided weights e.g. by only fine-tuning the classification head? Or do I have to train the Mask2Former model from scratch with my set of classes (including the partial classes + my new classes)?

[NazirNayal8]

Hi @pieris98,

1. yes, all models are trained with the 19 classes only.
2. You can of course fine-tune on new classes, but you will have to modify and retrain the region class probabilities head because it depends on the number of output classes, and everything else can stay the same. However I recommend also fine-tuning the object queries as well as the mask prediction head to get better results.