Video masked autoencoders (VideoMAE) are seen as data-efficient learners for self-supervised video pre-training (SSVP). Inspiration was drawn from the recent ImageMAE, and customized video tube masking with an extremely high ratio was proposed. Due to this simple design, video reconstruction is made a more challenging self-supervision task, leading to the extraction of more effective video representations during this pre-training process. Some hightlights of VideoMAE:

• Masked Video Modeling for Video Pre-Training
• A Simple, Efficient and Strong Baseline in SSVP
• High performance, but NO extra data required

This is a unofficial Keras implementation of VideoMAE: Masked Autoencoders are Data-Efficient Learners for Self-Supervised Video Pre-Training model. The official PyTorch implementation can be found here.

• [27-10-2023]: Code of Video-FocalNet in Keras becomes available.
• [15-10-2023]: Code of UniFormerV2 (UFV2) in Keras becomes available.
• [06-10-2023]: Code of Video Swin Transformer in Keras becomes available.
• [24-10-2023]: Kinetics-400 test data set can be found on kaggle, link.
• [9-10-2023]: TensorFlow SavedModel (formet) checkpoints, link.
• [6-10-2023]: VideoMAE integrated into Huggingface Space.
• [4-10-2023]: VideoMAE checkpoints SSV2 and UCF101 becomes available, link.
• [3-10-2023]: VideoMAE checkpoints on Kinetics-400 becomes available, link.
• [29-9-2023]: GPU(s), TPU-VM for fine-tune training are supported.
• [27-9-2023]: Code of VideoMAE in Keras becomes available.

git clone https://github.com/innat/VideoMAE.git
cd VideoMAE
pip install -e . 

There are many variants of VideoMAE mdoels available, i.e. small, base, large, and huge. And also for benchmark data specific, i.e. Kinetics-400, SSV2, and UCF101. Check this release and model zoo page to know details of it.

Only the inference part is provided for pre-trained VideoMAE models. Using the trained checkpoint, it would be possible to reconstruct the input sample even with high mask ratio. For end-to-end workflow, check this reconstruction.ipynb notebook. Some highlights:

from videomae import VideoMAE_ViTS16PT

# pre-trained self-supervised model
>>> model = VideoMAE_ViTS16PT(img_size=224, patch_size=16)
>>> model.load_weights('TFVideoMAE_B_K400_16x224_PT.h5')

# tube masking
>>> tube_mask = TubeMaskingGenerator(
    input_size=window_size, 
    mask_ratio=0.80
)
>>> make_bool = tube_mask()
>>> bool_masked_pos_tf = tf.constant(make_bool, dtype=tf.int32)
>>> bool_masked_pos_tf = tf.expand_dims(bool_masked_pos_tf, axis=0)
>>> bool_masked_pos_tf = tf.cast(bool_masked_pos_tf, tf.bool)

# running
>>> container = read_video('sample.mp4')
>>> frames = frame_sampling(container, num_frames=16)
>>> pred_tf = model(frames, bool_masked_pos_tf)
>>> pred_tf.numpy().shape
TensorShape([1, 1176, 1536])

A reconstructed results on a sample from SSV2 with mask_ratio=0.8

With the fine-tuned VideoMAE checkpoint, it would be possible to evaluate the benchmark datast and also retraining would be possible on custom dataset. For end-to-end workflow, check this quick retraining.ipynb notebook. It supports both multi-gpu and tpu-vm retraining and evaluation. Some highlights:

from videomae import VideoMAE_ViTS16FT

>>> model = VideoMAE_ViTS16FT(img_size=224, patch_size=16, num_classes=400)
>>> container = read_video('sample.mp4')
>>> frames = frame_sampling(container, num_frames=16)
>>> y = model(frames)
>>> y.shape
TensorShape([1, 400])

>>> probabilities = tf.nn.softmax(y_pred_tf)
>>> probabilities = probabilities.numpy().squeeze(0)
>>> confidences = {
    label_map_inv[i]: float(probabilities[i]) \
    for i in np.argsort(probabilities)[::-1]
}
>>> confidences

A classification results on a sample from Kinetics-400.

Video	Top-5
	{ 'playing_cello': 0.6552159786224365, 'snowkiting': 0.0018940207082778215, 'deadlifting': 0.0018381892004981637, 'playing_guitar': 0.001778001431375742, 'playing_recorder': 0.0017528659664094448 }

The pre-trained and fine-tuned models are listed in MODEL_ZOO.md. Following are some hightlights.

For Kinetrics-400, VideoMAE is trained around 1600 epoch without any extra data. The following checkpoints are available in both tensorflow SavedModel and h5 format.

^{?* Official ViT-H backbone of VideoMAE has weight issue in pretrained model, details MCG-NJU/VideoMAE#89. The FLOPs of encoder models (FT) are reported only.}

For SSv2, VideoMAE is trained around 2400 epoch without any extra data.

For UCF101, VideoMAE is trained around 3200 epoch without any extra data.

Some reconstructed video sample using VideoMAE with different mask ratio.

Kinetics-400-testset	mask
	0.8
	0.8
	0.9
	0.9

SSv2-testset	mask
	0.9
	0.9

UCF101-testset	mask
	0.8
	0.9

• Custom fine-tuning code.
• Publish on TF-Hub.
• Support Keras V3to support multi-framework backend.

If you use this videomae implementation in your research, please cite it using the metadata from our CITATION.cff file.

@inproceedings{tong2022videomae,
  title={Video{MAE}: Masked Autoencoders are Data-Efficient Learners for Self-Supervised Video Pre-Training},
  author={Zhan Tong and Yibing Song and Jue Wang and Limin Wang},
  booktitle={Advances in Neural Information Processing Systems},
  year={2022}
}