Introduction

If you replace the final ImageNet classifier with a new dense layer and train only on your custom classes, the model will stop predicting the original one thousand ImageNet labels. That is expected because you removed the head that produced those logits.

If you want both the original ImageNet classes and the new classes, you need a model design that preserves the old classifier or extends it. The core problem is not only architecture but also catastrophic forgetting during fine-tuning.

Why Replacing the Final Layer Removes ImageNet Classes

A typical transfer-learning workflow keeps the convolutional backbone and swaps the top classifier for a smaller task-specific layer. That is correct when you no longer need ImageNet predictions.

It is the wrong design when you still need the original labels. The final dense layer contains the mapping from extracted features to ImageNet class logits. Once you replace it, those logits no longer exist.

To preserve old classes, choose one of these patterns:

• keep the original ImageNet head and add a second head for new labels
• expand the classifier from 1000 outputs to 1000 + N outputs
• distill knowledge from the old model into the new one during training

A Two-Head Model Is the Safest Starting Point

A two-head model is simple and avoids disturbing the original classifier too early. One head stays frozen for ImageNet predictions, and the second head learns your new classes.

1import tensorflow as tf
2
3base = tf.keras.applications.MobileNetV2(
4    include_top=False,
5    weights="imagenet",
6    input_shape=(224, 224, 3),
7)
8base.trainable = False
9
10inputs = tf.keras.Input(shape=(224, 224, 3))
11x = tf.keras.applications.mobilenet_v2.preprocess_input(inputs)
12x = base(x, training=False)
13x = tf.keras.layers.GlobalAveragePooling2D()(x)
14
15imagenet_logits = tf.keras.layers.Dense(
16    1000,
17    name="imagenet_head",
18    trainable=False,
19)(x)
20new_logits = tf.keras.layers.Dense(5, name="custom_head")(x)
21
22model = tf.keras.Model(inputs, [imagenet_logits, new_logits])
23model.compile(
24    optimizer="adam",
25    loss={"custom_head": tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)},
26)

In practice, you usually load the original head weights rather than creating a random frozen dense layer as above. The point is structural: keep the old output path if you still need old predictions.

Expanding the Classifier Instead of Replacing It

If you want a single probability vector containing both old and new classes, expand the final layer rather than swapping it out. For example, turn a 1000-class head into a 1005-class head.

That requires training data and loss design that preserve the old classes. A common approach is to combine:

• normal classification loss on new labeled examples
• distillation loss that encourages the first one thousand logits to stay close to the original model

1student_logits = student(images, training=True)
2teacher_logits = teacher(images, training=False)
3
4classification_loss = tf.keras.losses.sparse_categorical_crossentropy(
5    new_targets,
6    student_logits[:, 1000:],
7    from_logits=True,
8)
9
10distillation_loss = tf.keras.losses.KLDivergence()(
11    tf.nn.softmax(teacher_logits / 2.0),
12    tf.nn.softmax(student_logits[:, :1000] / 2.0),
13)
14
15loss = tf.reduce_mean(classification_loss) + 0.5 * distillation_loss

This reduces forgetting because the student is asked to preserve the old behavior while learning the new labels.

Fine-Tuning Requires Restraint

Freezing the backbone first is usually the safest move. Once the new head learns reasonable features, unfreeze only the upper part of the backbone and continue with a small learning rate.

If you fine-tune the full network too aggressively on a narrow custom dataset, the feature extractor drifts away from the ImageNet representation. Even if you keep the old head, predictions on old classes can degrade badly.

That is why mixed training, rehearsal data, or distillation matters when old accuracy is important.

When You Should Not Combine Them

Sometimes the right answer is to keep two separate models. If your custom domain is very different from ImageNet or the deployment target is small, one combined model may add unnecessary complexity. In that case, route requests explicitly: use the ImageNet model for generic classification and the specialized model for domain-specific classes.

A simpler system with clear boundaries is often easier to debug and easier to deploy.

Common Pitfalls

• Replacing the final layer and expecting old ImageNet classes to remain available. That classifier was the part producing those labels.
• Fine-tuning the backbone too early. Small custom datasets can quickly erase useful pre-trained features.
• Ignoring catastrophic forgetting. Preserving architecture alone is not enough; training strategy matters too.
• Mixing old and new labels without a mapping plan. A single output layer needs a well-defined class index scheme.
• Assuming one combined model is always best. Sometimes two separate models are operationally cleaner.

Summary

• Swapping the final layer removes the original ImageNet output classes.
• To keep old classes, preserve the old head or expand it.
• A two-head architecture is the safest starting design.
• Distillation or rehearsal is important if you fine-tune while preserving old behavior.
• When complexity is not worth it, keep separate models for generic and custom classification.