Using Transformers to classify posture from video frames

Soham_Bhaumik: <p>I am fairly new to machine learning and I’ve been tinkering with transformers for a short while now.<br>
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:</p>
<pre><code class="lang-auto">  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)
        )
</code></pre>
<pre><code>    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
</code></pre>
<p>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.</p>
<p>Can anybody guide me with regards to this?<br>
Note: I have aroun 3500 training examples and 850 validation examples</p>