Implementation: Using simple fine-tuning (freezing all layers except the last few Linear layers) for LLMs

simple_fine_tuning.py

# Imports
from transformers import AutoModelForSequenceClassification

###################
model_uri = "distilbert/distilbert-base-uncased"
num_classes = 2

# Initialise the model
model = AutoModelForSequenceClassification.from_pretrained(
    model_uri, num_labels=num_classes
)

# Freeze all the layers
for param in model.parameters():
    param.requires_grad = False

# Unfreeze pre-classifier (penultimate layer)
for param in model.pre_classifier.parameters():
    param.requires_grad = True

# Unfreeze classifier (final layer)
for param in model.classifier.parameters():
    param.requires_grad = True