Video Classification Using SpatioTemporal Transformers
Background
In a previous series on transformers, we highlighted their advantages over RNNs and LSTMs for sequential data processing. Unlike sequential modelling in RNNs, transformers use attention mechanisms to capture long-range dependencies in a single pass, improving context retention and gradient flow via residual connections. Parallel processing also enables it in learning multiple attention contexts simultaneously, enhancing efficiency. This has led to architectures like BERT, enabling high-performance transfer learning across NLP, time-series forecasting, and vision tasks.
Speaking of vision tasks; the vision transformer architecture (ViT) improved image recognition performance with less computational complexity that state-of-the-art CNNs, becoming foundational for tasks like segmentation and detection. While BERTs and ViTs handle 1D sequences, video tasks require 2D spatiotemporal modelling. Building on the ViT architecture, a spatiotemporal transformer with sequential spatial and temporal attention outperforms state-of-the-art models for video classification like I3D and SlowFast while reducing training time.
ViT & Spatiotemporal Transformer
The ViT processes images by dividing, say a 224x224 image into 256 patches (16x 16) 0f 196 pixels, the embedding each 196-pixel patch. The resultant embeddings are combined with positional encodings to retain spatial information. A CLS token is appended to aggregate information for prediction before feeding the sequence into a transformer encoder for final prediction.
The spatiotemporal transformer classifies image sequences by processing, say a Tx224x224 input, dividing each frame into 256 patches (16x16) of 196 pixels and embedding them. This forms a Tx16x16xE representation, combined with positional encoding to retain spatial and temporal positions. It applies temporal attention to model patch evolution over time for each patch using a time encoder, then spatial attention between patches within each frame using a spatial encoder. The result is flattened to Tx256xE, with a CLS token added to aggregate learned structures for final prediction.
Get Emmanuel Olateju’s stories in your inbox
Join Medium for free to get updates from this writer.
Now that we understand how a spatiotemporal transformer works, let us implement one for the open source UCF50-Action-Recognition dataset.
Implementation Details
Rather than develop or spatiotemporal transformer from scratch, we make use of the huggingface TimesFormer package which implements the facebook spatiotemporal transformer described above and referred to as the TimesFormer.
Implementation
We begin by defining or clip parameters: the number of frames for classification and the height and width of frames, then we prepare our dataset to match this.
import os
import sys
import random
import math
import cv2
from sklearn.model_selection import StratifiedShuffleSplit
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, Subset, DataLoader
import torch.nn.functional as F
from einops import rearrange
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
T = SEQUENCE_LENGTH = 20
H, W = HEIGHT, WIDTH = 224, 224
ucf_path = "/kaggle/input/ucf50-action-recognition-dataset/UCF50"
actions = os.listdir(ucf_path)
N_CLASSES = 5
actions = ['PlayingTabla', 'HulaHoop', 'Lunges', 'Swing', 'Skijet']
# actions = random.sample(actions, N_CLASSES)
print(actions)
NC = 50
X = np.empty((0, T, H, W, 3))
y = []
for a, action in enumerate(actions):
N = X.shape[0]
clips = os.listdir(os.path.join(ucf_path, action))
for clip in tqdm(clips, desc=f"Extracting {action} Clip {a+1}/{len(actions)}", total=len(clips)):
path = os.path.join(ucf_path, action, clip)
cap = cv2.VideoCapture(path)
if not cap.isOpened():
continue
frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) # Total frames
frames = []
for f in range(frame_count):
ret, frame = cap.read()
if not ret:
break
frame = cv2.resize(frame, (HEIGHT, WIDTH))
frame = frame/255
frames.append(frame)
if len(frames)==T:
X = np.vstack((X, np.expand_dims(np.array(frames), axis=0)))
y.append(a)
frames = []
if X.shape[0] >= N + NC:
break
cap.release()
if X.shape[0] >= N + NC:
break
X = torch.tensor(X).to(torch.float32)
Y = torch.tensor(y)
X = rearrange(X, "b t h w c -> b t c h w")
time.sleep(0.25)
clear_output()
print(X.size(), Y.size())Now let us make a PyTorch dataset and dataloader
class UCF50Dataset(Dataset):
def __init__(self, X, y):
assert X.size(0) == len(y)
self.X = X
self.y = y
def __len__(self):
return self.X.size(0)
def __getitem__(self, idx):
if self.y != None:
return self.X[idx], self.y[idx]
return self.X[idx]
dataset = UCF50Dataset(X, Y)
train_batch_size = 5
val_batch_size = 3
print(train_batch_size, val_batch_size)
N = len(dataset)
split_ratio = 0.8
split_size = int(0.8*N)
split = StratifiedShuffleSplit(n_splits=1, test_size=1-split_ratio, random_state=42)
train_idx, val_idx = next(split.split(range(N), dataset.y))
train_dataset = Subset(dataset, train_idx)
val_dataset = Subset(dataset, val_idx)
train_loader = DataLoader(train_dataset, batch_size=train_batch_size, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=val_batch_size, shuffle=True)
print(f"Train samples: {len(train_dataset)}, Test samples: {len(val_dataset)}")Now let us define our TimesFormer hyperparameters
from transformers import TimesformerConfig, TimesformerModel, TimesformerForVideoClassification
from transformers import VivitConfig, VivitModel
N_PATCHES = 64
N_EMBD = 768
N_ATTN_HEADS = 6
N_LAYERS = 2
FF_INTERMEDIATE = 3072
INIT_STD = 0.2
timesformer_config_dict = {
"image_size": HEIGHT,
"patch_size": N_PATCHES,
"num_channels": 3,
"num_frames": T,
"hidden_size": N_EMBD,
"num_hidden_layers":N_LAYERS,
"num_attention_heads": N_ATTN_HEADS,
"intermediate_size": FF_INTERMEDIATE,
"hidden_act": "leaky_relu",
"hidden_dropout_prob": 0.05,
"attention_probs_dropout_prob": 0.1,
"initializer_range": INIT_STD,
"layer_norm_eps": 1E-6,
"qkv_bias": False,
"attention_type": "divided_space_time",
"drop_path_rate": 0.2,
}
timesformer_config = TimesformerConfig(**timesformer_config_dict)Now we can define our model which creates a TimesFormer instance, and then adds an extra layer computes the logits for softmax classification using the learned CLS token.
class ActionClassifier(nn.Module):
def __init__(self, n_embd, n_classes, TimesFormer):
super(ActionClassifier, self).__init__()
self.timesformer = TimesFormer
self.cls_layer = nn.Linear(N_EMBD, N_CLASSES)
def forward(self, x, y=None, device=DEVICE):
out = self.timesformer(x)
embeddings = out.last_hidden_state
cls_token = embeddings[:, 0, :]
logits = F.leaky_relu(self.cls_layer(cls_token))
if y==None:
return logits
elif isinstance(y, torch.Tensor):
loss = F.cross_entropy(logits, y, label_smoothing=0.1)
acc = accuracy_score(logits, y)
return logits, loss, acc
model = ActionClassifier(N_EMBD, N_CLASSES, TimesformerModel(timesformer_config))
model =model.to(DEVICE)
def init_weights(m):
if isinstance(m, nn.Linear): # Apply only to Linear layers
torch.nn.init.normal_(m.weight, mean=0.0, std=INIT_STD)
if m.bias is not None:
torch.nn.init.zeros_(m.bias) # Set bias to 0
model.apply(init_weights)Then we can train our model
torch.manual_seed(42)
np.random.seed(42)
n_iter = 3000
eval_inter = 5
n_evals = 3
lr_ = 3E-4
betas_ = (0.9, 0.999)
optimizer = optim.Adam(model.parameters(), lr=lr_, betas=betas_)
lr_scheduler = LRScheduler(lr_, patience=5, factor=1, margin=0.05, n_evals=n_iter//n_evals)
early_stop = EarlyStopping(patience=10, margin=0.1)
chk_pt = CheckpointModel(margin=0.05)
history = {
"train_loss": [],
"val_loss": [],
"train_acc": [],
"val_acc": [],
"iter": [],
"lr": [],
}
for i in tqdm(range(n_iter), total=n_iter):
if i==0 or i%eval_inter==0 or i>=n_iter-1:
model = model.eval()
with torch.no_grad():
train_loss=torch.zeros((n_evals,))
train_acc=torch.zeros((n_evals,))
for j in range(n_evals):
with torch.autocast(device_type=DEVICE):
x, y = next(iter(train_loader))
x = x.to(DEVICE)
y = y.to(DEVICE)
logits, loss, acc = model(x, y, device=DEVICE)
train_loss[j] = loss
train_acc[j] = acc
val_loss = torch.zeros((n_evals,))
val_acc = torch.zeros((n_evals,))
for j in range(n_evals):
with torch.autocast(device_type=DEVICE):
x, y = next(iter(val_loader))
x = x.to(DEVICE)
y = y.to(DEVICE)
logits, loss, acc = model(x, y, device=DEVICE)
val_loss[j] = loss
val_acc[j] = acc
train_loss = train_loss.mean().item()
train_acc = train_acc.mean()
val_loss = val_loss.mean().item()
val_acc = val_acc.mean()
# lr_ = lr_scheduler.exponential_lr(val_loss)
# for param_group in optimizer.param_groups:
# param_group["lr"] = lr_
history["train_loss"].append(train_loss)
history["val_loss"].append(val_loss)
history["train_acc"].append(train_acc)
history["val_acc"].append(val_acc)
history["iter"].append(i)
history["lr"].append(lr_)
early_stop_signal = early_stop.check(train_loss, val_loss)
if early_stop_signal == True:
break
chk_pt.checkpoint(val_loss, model)
model = model.train()
optimizer.zero_grad()
with torch.autocast(device_type=DEVICE):
x, y = next(iter(train_loader))
x, y = x.to(DEVICE), y.to(DEVICE)
logits, loss, acc = model(x, y, device=DEVICE)
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step() After training, let us see our training loss, accuracy, and confusion matrix, which tells us how well our model performs on the validation set.
Real Time Inference
Let us run inference on some clips with the resulting video shown
model.load_state_dict(chk_pt.best_state)
model.eval()
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
output_video_path = "continuous_inference.mp4"
out = None
for a, action in enumerate(actions):
clips = os.listdir(os.path.join(ucf_path, action))[:5]
for clip in tqdm(clips, desc=action):
path = os.path.join(ucf_path, action, clip)
cap = cv2.VideoCapture(path)
if not cap.isOpened():
print(f"Skipping file: {path}")
continue
frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
fps = int(cap.get(cv2.CAP_PROP_FPS))
if out is None:
out = cv2.VideoWriter(output_video_path, fourcc, fps, (width, height))
frames = []
for f in range(frame_count):
ret, frame = cap.read()
if not ret:
break
frame_resized = cv2.resize(frame, (HEIGHT, WIDTH))
frame_resized = frame_resized/255
if len(frames) == T:
frames = frames[1:]
frames.append(frame_resized)
if len(frames)==T:
X = torch.tensor(np.array(frames), dtype=torch.float32).to(DEVICE)
X = X.unsqueeze(dim=0)
X = rearrange(X, "b t h w c -> b t c h w")
logits = model(X)
probs = F.softmax(logits, dim=-1)
pred = torch.argmax(probs, dim=-1)[0].item()
pred = actions[pred]
font = cv2.FONT_HERSHEY_SIMPLEX
cv2.putText(frame, f"Prediction: {pred}", (20,50), font, 0.5, (0, 255, 0), 2, cv2.LINE_AA)
out.write(frame)
cap.release()
if out is not None:
out.release()
print(f"Saved classified video: {output_video_path}")Conclusion
In this write-up, we explored spatiotemporal transformers for video classification, building on the foundational ViT architecture. We implemented a TimesFormer model using Hugging Face’s package and applied it to the UCF50-Action-Recognition dataset. Our approach involved preprocessing video clips, defining a custom dataset, and training a transformer-based model capable of capturing both spatial and temporal dependencies.
The results highlight the potential of spatiotemporal transformers in video classification, potentially outperforming traditional CNN-based models while maintaining computational efficiency. Our validation metrics and real-time inference demonstrate the model’s ability to generalize well across diverse action categories.
