[Python] Trying to classify sequences of video frames as a particular posture for dogs

I am fairly new to machine learning and I've been tinkering with transformers for a short while now. I have written a transformer architecture that should in my opinion be able to understand why this architecture does not work. I am trying to fuse modalities from a vit- based dog breed classifier and a vit base canine pose estimator using the following code:

import torch
import torch.nn as nn
from transformers import AutoImageProcessor, AutoModelForImageClassification

processor = AutoImageProcessor.from_pretrained("skyau/dog-breed-classifier-vit")
backbone_model = AutoModelForImageClassification.from_pretrained("skyau/dog-breed-classifier-vit")
class LearnablePositionalEncoding(nn.Module):
def __init__(self, seq_len, embed_dim):
super().__init__()
self.pos_embed = nn.Parameter(torch.zeros(1, seq_len, embed_dim)) # (1, T, D)
nn.init.trunc_normal_(self.pos_embed, std=0.02) # or use normal_

def forward(self, x):
x = x + self.pos_embed[:, :x.size(1), :]
return x

class CaninePostureClassifier(nn.Module):
def __init__(self, num_keypoints = 20, keypoint_dim = 2, hidden_dim = 256, embed_dim = 128, num_classes = 4):
super().__init__()
self.embed_dim = embed_dim
self.hidden_dim = hidden_dim
self.keypoint_dim = keypoint_dim
self.num_keypoints = num_keypoints
self.backbone = backbone_model
self.kpt_projections = nn.Sequential(
nn.Linear(keypoint_dim, self.hidden_dim),
nn.ReLU(),
nn.Dropout(0.1),
nn.LayerNorm(self.hidden_dim),
nn.Linear(hidden_dim, self.embed_dim),
nn.ReLU(),
nn.LayerNorm(embed_dim),
nn.Dropout(0.1)
)


self.image_size = (224, 224)

self.projection_out = nn.Sequential(
nn.Linear(self.embed_dim, self.backbone.config.hidden_size),
nn.Dropout(0.1)
)
self.temporal_pos_enc = LearnablePositionalEncoding(seq_len=16, embed_dim=self.embed_dim)
self.joint_encoder_layer = nn.TransformerEncoderLayer(d_model=self.embed_dim, nhead=4, dropout=0.2, batch_first=True)
self.kp_self_attn = nn.TransformerEncoder(self.joint_encoder_layer, num_layers=2)

self.temporal_attn_layer = nn.TransformerEncoderLayer(d_model=self.embed_dim, nhead=4, dropout=0.2, batch_first=True)
self.temporal_self_attn = nn.TransformerEncoder(self.temporal_attn_layer, num_layers=2)


self.temporal_patch_attn = nn.TransformerEncoder(
nn.TransformerEncoderLayer(d_model=768, nhead=4, dropout=0.1),
num_layers=1
)

for param in self.backbone.parameters():
param.requires_grad = False

self.classifier = nn.Sequential(
nn.LayerNorm(self.backbone.config.hidden_size),
nn.Dropout(0.2),
nn.Linear(self.backbone.config.hidden_size , 4)
)



def forward(self, rgb, keypoints):
#perform attention spatially and temporally over the input keypoint sequence
B, T, K, D = keypoints.shape

# keypoints = keypoints.clone()
# mean = keypoints.mean(dim=1, keepdim=True)
# std = keypoints.std(dim=1, keepdim=True) + 1e-6
# keypoints = (keypoints - mean) / std # (B, T, K, D)

kpt_embed = self.kpt_projections(keypoints)
jointwise_in = kpt_embed.reshape(-1, self.num_keypoints, self.embed_dim)
jointwise_out = self.kp_self_attn(jointwise_in)
jointwise_out = jointwise_out.view(B, T, self.num_keypoints, self.embed_dim)
jointwise_out = jointwise_out.permute(0, 2, 1, 3)
temporal_in = jointwise_out.reshape(-1, T, self.embed_dim)
temporal_in = self.temporal_pos_enc(temporal_in)
temporal_out = self.temporal_self_attn(temporal_in)
temporal_out = temporal_out.view(B, T * K, self.embed_dim)
keypoint_proj_out = self.projection_out(temporal_out)
keypoint_proj_out = keypoint_proj_out.view(B, T, K, self.backbone.config.hidden_size)
keypoint_tokens = keypoint_proj_out.reshape(B, T * K, self.backbone.config.hidden_size)


#extract features from a vitpose backbone
patch_tokens = []

for t in range(T):
frame = rgb[:, t]
with torch.no_grad():
token = self.backbone.vit(frame).last_hidden_state[:,1: ,:]
patch_tokens.append(token)

patch_tokens = torch.stack(patch_tokens)
patch_tokens = patch_tokens.permute(0, 2, 1, 3) # [B, 196, T, 768]
patch_tokens = patch_tokens.reshape(-1, T, 768) # Treat each spatial location across time as a temporal sequence
patch_tokens = self.temporal_patch_attn(patch_tokens) # your nn.TransformerEncoder
patch_tokens = patch_tokens.view(B, 196, T, 768).permute(0, 2, 1, 3) # [B, T, 196, 768]
patch_tokens = patch_tokens.reshape(B, T * 196, 768) # ready to concat with keypoint tokens

fused = torch.cat([patch_tokens, keypoint_tokens], dim=1)
# pooled = AttnWeightedPooling(hidden_dim = self.backbone.config.hidden_size)(fused_tokens)
pooled = fused.mean(dim = 1)
out = self.classifier(pooled)
return out


when I use cross attention between the modalities, it performs decently well (macro F1 of 74% and val accuracy of 78.5%). but I really want to understand why this model does not perform well at all.what i have been trying to o, model how my ground truth keypoints eveolve spatially and then temporally, by splitting them into appropriate tokens, and the fusing that with rgb features before averaging the token values and classifying. this model tops out at 55% val accuracy.

Can anybody guide me with regards to this? Note: I have aroun 3500 training examples and 850 validation examples

Continue reading...