To use a pre-trained TimeSformer model for inference, import the TimeSformer class from timesformer.models.vit. You must specify the img_size, num_classes, num_frames, and attention_type. To load weights, provide the path to a .pyth model file via the pretrained_model argument.
Input tensors should be in the shape (batch, channels, frames, height, width).
from pathlib import Path
import torch
from timesformer.models.vit import TimeSformer
# Define path to pre-trained weights
model_file = Path.home()/'TimeSformer/models/TimeSformer_divST_8x32_224_K600.pyth'
# Initialize model
model = TimeSformer(
img_size=224,
num_classes=600,
num_frames=8,
attention_type='divided_space_time',
pretrained_model=str(model_file)
)
# Prepare dummy video: (batch x channels x frames x height x width)
dummy_video = torch.randn(2, 3, 8, 224, 224)
# Perform inference
pred = model(dummy_video)
# Output shape will be (batch, num_classes)
assert pred.shape == (2, 600)