Remove layer in pretrained model

Beginners
1

I’m new to ML, I’m trying to perform an ablation study. Is there a concise way to access and remove layers from a pretrained model? Specifically, I’m working on this code:

pretrained_model_name = "nvidia/mit-b0"
model = AutoModelForSemanticSegmentation.from_pretrained(
    pretrained_model_name, id2label=id2label, label2id=label2id
)

# to-do: something like "model.remove_layers(-1)"

I know there’s model.config, but it doesn’t seem to have what I’m looking for. Any tips? Thanks!

1 Like
2

This worked for me :slight_smile:

model = AutoModelForSequenceClassification.from_pretrained(“BAAI/bge-small-en-v1.5”, num_labels=3, id2label=id2label, label2id=label2id)

model.base_model.encoder.layer = model.base_model.encoder.layer[:-1]

1 Like