TransformerEncoder.py

from torch import nn
import torch

class MLP(nn.Module):
    def __init__(self,embed_dim,expansion,dropout):
        super().__init__()
        self.fc1=nn.Linear(
            in_features=embed_dim,
            out_features=embed_dim*expansion
        )
        self.gelu=nn.GELU()
        self.dropout1=nn.Dropout(dropout)
        self.fc2=nn.Linear(
            in_features=embed_dim*expansion,
            out_features=embed_dim
        )
        self.dropout2=nn.Dropout(dropout)

    def forward(self,x):
        x=self.fc1(x)
        x=self.gelu(x)
        x=self.dropout1(x)
        x=self.fc2(x)
        x=self.dropout2(x)
        return x
    

class TransformerEncoderLayer(nn.Module):
    def __init__(self,embed_dim,num_heads,expansion,dropout):
        super().__init__()
        self.layernorm1=nn.LayerNorm(embed_dim)
        self.MHA=nn.MultiheadAttention(
            embed_dim=embed_dim,
            num_heads=num_heads,
            dropout=dropout,
            batch_first=True)
        self.dropout=nn.Dropout(dropout)
        self.layernorm2=nn.LayerNorm(embed_dim)
        self.mlp=MLP(
            embed_dim=embed_dim,
            expansion=expansion,
            dropout=dropout
        )

    def forward(self,x):
        residual1=x

        x=self.layernorm1(x)
        x=self.MHA(x,x,x)[0]
        x=x+residual1

        residual2=x

        x=self.layernorm2(x)
        x=self.mlp(x)
        x=x+residual2

        return x
    

#---------------------------------------------------------------------------------------------------------------
if __name__ == "__main__":
    # Test parameters
    BATCH_SIZE = 4
    NUM_PATCHES = 196          # e.g., 224x224 image with 16x16 patches → 196 + 1 cls token
    EMBED_DIM = 768
    NUM_HEADS = 12
    EXPANSION = 4
    DROPOUT = 0.1

    # Create dummy input (with class token already included)
    x = torch.randn(BATCH_SIZE, NUM_PATCHES + 1, EMBED_DIM)
    print(f"Input shape : {x.shape}")

    # Initialize layer
    encoder_layer = TransformerEncoderLayer(
        embed_dim=EMBED_DIM,
        num_heads=NUM_HEADS,
        expansion=EXPANSION,
        dropout=DROPOUT
    )

    # Forward pass
    output = encoder_layer(x)

    print(f"Output shape: {output.shape}")
    print(f"Expected shape: ({BATCH_SIZE}, {NUM_PATCHES + 1}, {EMBED_DIM})")

    # Assertions
    assert output.shape == x.shape, "❌ Shape mismatch!"
    assert not torch.isnan(output).any(), "❌ NaN values detected!"
    
    print("✅ TransformerEncoderLayer test passed successfully!\n")

    # Optional: Check parameter count
    total_params = sum(p.numel() for p in encoder_layer.parameters() if p.requires_grad)
    print(f"Total trainable parameters in one layer: {total_params:,}")