How to Modify UperNetForSemanticSegmentation from 150 Classes to Binary Classes While Retaining Pre-Trained Weights

Models
1

Hello everyone,

I’m working with the UperNetForSemanticSegmentation model using a Swin backbone, specifically loading it with pre-trained weights from "openmmlab/upernet-swin-tiny". The current configuration of the model is set up to predict 150 classes, but I would like to modify it to perform binary segmentation (i.e., 2 classes).

Here’s the code snippet I have so far:

from transformers import UperNetForSemanticSegmentation, SwinConfig, UperNetConfig

backbone_config = SwinConfig()

config = UperNetConfig(backbone_config=backbone_config)

model = UperNetForSemanticSegmentation(config).from_pretrained("openmmlab/upernet-swin-tiny")

Could someone guide me on how to change the number of output classes to 2 while retaining the pre-trained weights for the rest of the model? Any specific steps or code examples would be greatly appreciated.

Thank you!