How can we load the checkpoint from movinet to fine tune?

[thuanbui1309]

Hi, I am trying to fine tuning the pre-trained movinet a2 stream model with a subset of ucf101 dataset. Originally the movinet a2 has 600 classes as output, but the subset I want to fine tune on only has 10 classess. I try to load the check points and training from that but it keeps saying mismatch in num classess. Below is my code:
`
model_id = 'a2'
use_positional_encoding = model_id in {'a3', 'a4', 'a5'}
resolution = 224
backbone = movinet.Movinet(
model_id=model_id,
causal=True,
conv_type='2plus1d',
se_type='2plus3d',
activation='hard_swish',
gating_activation='hard_sigmoid',
use_positional_encoding=use_positional_encoding,
use_external_states=False,
)

def build_classifier(batch_size, num_frames, resolution, backbone, num_classes):
"""Builds a classifier on top of a backbone model."""
model = movinet_model.MovinetClassifier(
backbone=backbone,
num_classes=num_classes)
model.build([batch_size, num_frames, resolution, resolution, 3])

return model

Construct loss, optimizer and compile the model

with distribution_strategy.scope():
model = build_classifier(batch_size, num_frames, resolution, backbone, 10)
loss_obj = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
optimizer = tf.keras.optimizers.Adam(learning_rate = 0.001)
model.compile(loss=loss_obj, optimizer="Adam", metrics=['accuracy'])
`

[nmahammad]

The fix is the standard MoViNet transfer-learning pattern: load into a 600-class classifier first, then build a fresh 10-class classifier on that same backbone object. The backbone weights carry over by reference; only the head is new and random.

`
backbone = movinet.Movinet(
model_id='a2',
causal=True,
conv_type='2plus1d',
se_type='2plus3d',
activation='hard_swish',
gating_activation='hard_sigmoid',
use_positional_encoding=False,
use_external_states=False,
)

1) Build a 600-class classifier so it MATCHES the checkpoint's head.

model = movinet_model.MovinetClassifier(backbone=backbone, num_classes=600)
model.build([None, None, None, None, 3])

2) Restore the checkpoint into the 600-class model to fill the backbone.

checkpoint_path = tf.train.latest_checkpoint(checkpoint_dir)
checkpoint = tf.train.Checkpoint(model=model)
status = checkpoint.restore(checkpoint_path)
status.assert_existing_objects_matched() # confirms it actually loaded

3) NOW build the 10-class classifier on the SAME backbone (weights persist).

model = build_classifier(batch_size, num_frames, resolution, backbone, 10)

`

You can't load 600-class weights into a 10-class head, so you load at 600 classes, and only after the backbone is loaded, you swap 10-class head. Because backbone is the same Python object passed to both classifiers, the restored weights stay the same and the new head is initialized randomly and trained.

[UdaySharmaGitHub]

Hi @thuanbui1309 ,

The issue is that the pretrained MoViNet A2 checkpoint was created with a 600-class classification head, while you're building the classifier with num_classes=10.

You generally don't want to load the 600-class checkpoint directly into a 10-class classifier because the final classification layer has a different shape.

The easiest approach is to restore the pretrained checkpoint into a 600-class classifier first, and then create a new 10-class classifier using the same backbone.

For example:

model_id = "a2"

backbone = movinet.Movinet(
    model_id=model_id,
    causal=True,
    conv_type="2plus1d",
    se_type="2plus3d",
    activation="hard_swish",
    gating_activation="hard_sigmoid",
    use_positional_encoding=False,
    use_external_states=False,
)

# 1. Create the classifier with the SAME number of classes
#    as the pretrained checkpoint.
model_600 = movinet_model.MovinetClassifier(
    backbone=backbone,
    num_classes=600,
)

model_600.build([None, None, 224, 224, 3])

# 2. Restore the pretrained checkpoint.
checkpoint_path = tf.train.latest_checkpoint(checkpoint_dir)

checkpoint = tf.train.Checkpoint(model=model_600)
status = checkpoint.restore(checkpoint_path)

status.assert_existing_objects_matched()

# 3. Create a new classifier with your 10 classes
#    using the SAME backbone.
model_10 = movinet_model.MovinetClassifier(
    backbone=backbone,
    num_classes=10,
)

model_10.build([batch_size, num_frames, 224, 224, 3])

# 4. Compile and fine-tune.
loss_obj = tf.keras.losses.SparseCategoricalCrossentropy(
    from_logits=True
)

optimizer = tf.keras.optimizers.Adam(
    learning_rate=0.001
)

model_10.compile(
    loss=loss_obj,
    optimizer=optimizer,
    metrics=["accuracy"],
)

Why this works

The important part is that both classifiers use the same backbone object:

model_600 = MovinetClassifier(
    backbone=backbone,
    num_classes=600,
)

After restoring the checkpoint, the pretrained backbone weights are loaded into that backbone.

Then:

model_10 = MovinetClassifier(
    backbone=backbone,
    num_classes=10,
)

creates a new 10-class classification head on top of the already-loaded backbone.

So conceptually you have:

Pretrained checkpoint
        │
        ▼
600-class classifier
        │
        ├── restore pretrained backbone weights
        │
        ▼
   MoViNet A2 backbone
        │
        ▼
new 10-class classification head
        │
        ▼
   Your UCF101 subset

The 600-class classification head from the original checkpoint is not reused. It is only used to make the checkpoint structure compatible during restoration.

Your new 10-class head starts with newly initialized weights and is then trained on your 10-class dataset.

One thing to watch for

I would also recommend checking the checkpoint restoration carefully rather than using expect_partial() and ignoring missing variables.

For example:

status = checkpoint.restore(checkpoint_path)
status.assert_existing_objects_matched()

This helps verify that the expected pretrained variables were actually restored.

Also make sure the preprocessing used for your UCF101 videos matches the preprocessing expected by the MoViNet model, particularly the frame shape/order and normalization.

If you want to fine-tune only the new classification head initially, you can also freeze the backbone:

backbone.trainable = False

Train the 10-class head first, then unfreeze the backbone and continue with a smaller learning rate for full fine-tuning.

This is the standard transfer-learning pattern here: checkpoint-compatible model → restore pretrained backbone → replace classification head → fine-tune on the new class set.

GitHub: https://github.com/UdaySharmaGithub
If this answer was helpful, feel free to check out my profile for other projects and contributions.
Hope this helps!