Re-train model with new classes
ML System Design practice on Codemia
Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.
Introduction
Adding new classes is not a minor update. In most classification systems, the output layer, label mapping, metrics, and evaluation set all depend on the class vocabulary. When that vocabulary changes, you usually need more than "one more training run."
The main design question is whether you want a true incremental learner or a standard model that is being expanded to predict additional classes. Most deep-learning classifiers fall into the second category.
What Changes When New Classes Appear
Suppose a model originally predicts two classes, cat and dog. If you now want rabbit as well, several artifacts change at once:
- the final classifier layer size
- the label-to-index mapping
- the training set and validation set
- any business logic built around class probabilities or thresholds
That is why simply calling fit on new rabbit examples is usually not enough. The old model literally does not have an output neuron for the new class.
The Safe Default: Retrain With Old and New Data
The most reliable approach is to train a new classifier over the full class set using a dataset that includes both old and new classes. This avoids catastrophic forgetting, where the model learns the new class and degrades badly on the old ones.
If you have a useful existing backbone, you do not need to start from random weights. Reuse feature-extracting layers, replace the classifier head, and retrain on the expanded label space.
A Keras Example
This small example trains a two-class model, then replaces the output layer with a three-class head:
Now expand the classifier:
This is the transfer-learning pattern: keep the learned representation, replace the head, and train on the new class set.
Why Mixing Old and New Examples Matters
If x_new contains mostly rabbit data and very little cat or dog data, the classifier can drift toward the new distribution. Even if the feature extractor is still useful, the new head only learns what you show it.
A better dataset includes:
- enough examples from old classes to preserve decision boundaries
- enough examples from the new classes to learn separation
- validation splits that measure both old-class and new-class performance
When old data is unavailable, class expansion becomes much harder. You may need distillation, rehearsal buffers, or specialized continual-learning methods. Those approaches are real research areas because the problem is genuinely difficult.
Fine-Tuning Versus Full Retraining
Use fine-tuning when:
- the new classes are related to the old problem
- the old backbone still extracts useful features
- you have limited compute or training time
Use full retraining when:
- the class set changed substantially
- the data distribution changed, not just the labels
- the old model was already near its limits
In Keras-style workflows, a common sequence is:
- freeze the old backbone
- train the new head
- unfreeze some or all of the backbone
- continue training with a lower learning rate
That matches the standard transfer-learning guidance.
Update the Label Vocabulary Too
Do not forget the metadata. If the original model used:
the expanded model must use:
This sounds obvious, but many deployment bugs happen here. The numeric model is updated, but the serving code still decodes outputs with the old class list.
What About Incremental Learning?
Some scikit-learn estimators support partial_fit, and some continual-learning methods try to absorb new classes over time. Those are useful when retraining on full historical data is impossible.
But for ordinary neural-network classifiers, class expansion is still usually handled by rebuilding the classifier head and retraining with a mixed dataset. That is the baseline approach you should expect to implement first.
Common Pitfalls
The biggest mistake is training only on the new class and expecting the model to retain old knowledge. That usually causes catastrophic forgetting.
Another is forgetting to resize the final layer. A two-class output cannot produce a third class no matter how much new data you add.
Teams also overlook label metadata, confusion-matrix design, and monitoring dashboards. All of those need the new class list too.
Finally, evaluate per class. A single accuracy number can hide the fact that the new class improved while the old ones collapsed.
Summary
- Adding new classes usually requires changing the output layer and label mapping.
- The safest baseline is retraining on a dataset that contains both old and new classes.
- Reusing the old backbone and replacing the classifier head is a common transfer-learning pattern.
- Mixing old and new examples reduces catastrophic forgetting.
- Update deployment metadata and evaluate each class separately after expansion.
Related reading
- Read big train/validation/test datasets in tensorflow
- Read csv files in a MLFlow pipeline
- Read in Large CSV File and feed into TensorFlow
- Read mnist images into Tensorflow
- Read .mat files in Python
- Reading a file from a private S3 bucket to a pandas dataframe
- Read only mode in keras
- Reading data from bucket in Google ml-engine tensorflow
.png&w=3840&q=75)
Tackling System Design Interview Problems
A short course that equips you with the skills to approach system design interviews methodically.
Start the free courseTrack what you have practised
A free account saves your progress, solutions and study plan across every problem on Codemia.
ML System Design practice on Codemia
Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.