kfkas/movinet-a0-stream-pytorch

MoViNet-A0 Stream PyTorch

Trainable PyTorch and Transformers implementation of causal MoViNet-A0 Stream. The checkpoint contains the Kinetics-600 weights ported by Atze00/MoViNet-pytorch from the official Google Research MoViNet release.

The model accepts clips in Transformers video layout:

[batch, frames, channels, height, width]

The reference spatial resolution is 172x172, and RGB values should be floating point values in [0, 1].

pip install "torch>=2.1" "transformers>=5.0" \
  "fvcore>=0.1.5.post20210630" "einops>=0.8"

Load and train

import torch
from transformers import AutoModelForVideoClassification

model = AutoModelForVideoClassification.from_pretrained(
    "kfkas/movinet-a0-stream-pytorch",
    revision="0c1fcb4911f4971afc984bca813af2f9c9ab9cdd",
    trust_remote_code=True,
)

pixel_values = torch.rand(2, 8, 3, 172, 172)
labels = torch.tensor([10, 25])
output = model(pixel_values=pixel_values, labels=labels)
output.loss.backward()

For a new task, replace the Kinetics classifier:

model = AutoModelForVideoClassification.from_pretrained(
    "kfkas/movinet-a0-stream-pytorch",
    revision="0c1fcb4911f4971afc984bca813af2f9c9ab9cdd",
    num_labels=3,
    ignore_mismatched_sizes=True,
    trust_remote_code=True,
)

Frame-by-frame streaming

Independent clip calls reset causal state automatically. To retain state across successive frames, reset once and pass use_stream_state=True:

model.reset_stream()
for frame in frames:
    output = model(
        pixel_values=frame[:, None],
        use_stream_state=True,
    )

frame has shape [batch, channels, height, width]. Call reset_stream() before starting another video.

Provenance

• Architecture: MoViNets: Mobile Video Networks for Efficient Video Recognition
• PyTorch port and checkpoint: Atze00/MoViNet-pytorch, commit c2d1edf48fc6c5259707f9d833f22171b4f63493
• Parameters: 3,748,031
• Classes: 600
• License: MIT

This repository uses custom Transformers code. Review the modeling files and pin a repository revision for production use.