model_builder.py

from torch import nn
from torchvision import models


class Baseline(nn.Module):
    def __init__(self, input_shape: int, hidden_units: int, output_shape: int):
        super(Baseline, self).__init__()
        self.conv_layer = nn.Sequential(
            nn.Conv2d(input_shape, hidden_units, kernel_size=3),
            nn.ReLU(),
            nn.Conv2d(hidden_units, hidden_units, kernel_size=3),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )
        self.feed_forward = nn.Sequential(
            nn.Flatten(),
            nn.Linear(hidden_units * 30 ** 2, output_shape)
        )

    def forward(self, inputs):
        return self.feed_forward(self.conv_layer(inputs))


class TinyVGG(nn.Module):
    def __init__(self, input_shape: int, hidden_units: int, output_shape: int):
        super(TinyVGG, self).__init__()
        self.conv_layer1 = nn.Sequential(
            nn.Conv2d(input_shape, hidden_units, kernel_size=3),
            nn.ReLU(),
            nn.Conv2d(hidden_units, hidden_units, kernel_size=3),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )
        self.conv_layer2 = nn.Sequential(
            nn.Conv2d(hidden_units, hidden_units, kernel_size=3),
            nn.ReLU(),
            nn.Conv2d(hidden_units, hidden_units, kernel_size=3),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )
        self.feed_forward = nn.Sequential(
            nn.Flatten(),
            nn.Linear(hidden_units * 13 ** 2, output_shape)
        )
    def forward(self, inputs):
        return self.feed_forward(self.conv_layer2(self.conv_layer1(inputs)))


class EfficientNet(nn.Module):
    def __init__(self, *args, **kwargs):
        super(EfficientNet, self).__init__()

        # Load the EfficientNet-B0 model with pre-trained weights
        efficientnet_model = models.efficientnet_b0(weights=models.EfficientNet_B0_Weights.DEFAULT)

        # Freeze the parameters in the feature extractor (backbone)
        for param in efficientnet_model.parameters():
            param.requires_grad = False

        # Replace the classifier head (the final layer) with your custom head for binary classification
        self.efficientnet_features = efficientnet_model.features
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(1280, 2)  # Binary classification (2 output classes)
        )

    def forward(self, inputs):
        return self.classifier(self.efficientnet_features(inputs))


'''
Below are code for Goog;es Vision Transformer
'''

class PatchEmbeddingLayer(nn.Module):
    def __init__(self, in_channels, embedding_dim, patch_size):
        super().__init__()
        self.patch_size = patch_size
        self.conv_layer = nn.Conv2d(in_channels=in_channels,
                                   out_channels=embedding_dim,
                                   kernel_size=patch_size,
                                   stride=patch_size)
        self.flatten = nn.Flatten(start_dim=2, end_dim=3)

    def forward(self, x):
        return self.flatten(self.conv_layer(x)).permute(0, 2, 1)


class MultiHeadSelfAttentitionBlock(nn.Module):
    def __init__(self, embedding_dim, num_head, dropout:float=0.1):
        super().__init__()
        self.norm_layer = nn.LayerNorm(embedding_dim)
        self.multi_atten = nn.MultiheadAttention(embed_dim=embedding_dim,
                                                num_heads=num_head,
                                                dropout=dropout)

    def forward(self, x):
        x = self.norm_layer(x)
        atten_output, _ = self.multi_atten(x, x, x, need_weights=False)
        return atten_output


class MultiLayerPerceptronBlock(nn.Module):
    def __init__(self, in_features, embedding_dim, dropout=0.1):
        super().__init__()
        self.layer_norm = nn.LayerNorm(embedding_dim)
        self.mlp_layer = nn.Sequential(
            nn.Linear(in_features=embedding_dim,
                      out_features=in_features),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(in_features=in_features,
                     out_features=embedding_dim)
        )

    def forward(self, x):
        return self.mlp_layer(self.layer_norm(x))


class TransformerEncoderBlock(nn.Module):
    def __init__(self, embedding_dim, num_head, mlp_hidden_units, msa_dropout=0.1, mlp_dropout=0.1):
        super().__init__()
        self.msa = MultiHeadSelfAttentitionBlock(embedding_dim,
                                                num_head,
                                                msa_dropout)
        self.mlp = MultiLayerPerceptronBlock(mlp_hidden_units,
                                            embedding_dim,
                                            mlp_dropout)

    def forward(self, x):
        x = self.msa(x) + x
        x = self.mlp(x) + x
        return x


class VisionTransformer(nn.Module):
    def __init__(self,
                 input_shape=3,
                 image_size=224,
                 patch_size=16,
                 msa_head=12,
                 msa_dropout=0,
                 mlp_dropout=0.1,
                 embedding_dropout=0.1,
                 mlp_hidden_units=3072,
                 num_transformer_block=12,
                 num_classes=2):
        super().__init__()
        self.embedding_dim = int((patch_size ** 2) * 3)
        self.num_patches = int((image_size ** 2) / (patch_size ** 2))
        self.patcher = PatchEmbeddingLayer(3, self.embedding_dim, patch_size)
        self.class_token_embedding = nn.Parameter(
            torch.randn(1, 1, self.embedding_dim)
        )
        self.position_eembedding = nn.Parameter(
            torch.randn(1, self.num_patches + 1, self.embedding_dim)
        )
        self.embedding_dropout = nn.Dropout(embedding_dropout)
        self.transformer_encoder = nn.Sequential(*[
            TransformerEncoderBlock(self.embedding_dim,
                                   msa_head,
                                   mlp_hidden_units,
                                   msa_dropout,
                                   mlp_dropout)
            for _ in range(num_transformer_block)
        ])
        self.classifer_layer = nn.Sequential(
            nn.LayerNorm(self.embedding_dim),
            nn.Linear(in_features=self.embedding_dim,
                     out_features=num_classes)
        )

    def forward(self, x):
        # 12. Get batch size
        batch_size = x.shape[0]

        # 13. Create class token embedding and expand it to match the batch size (equation 1)
        class_token = self.class_token_embedding.expand(batch_size, -1, -1) # "-1" means to infer the dimension (try this line on its own)

        # 14. Create patch embedding (equation 1)
        x = self.patcher(x)

        # 15. Concat class embedding and patch embedding (equation 1)
        x = torch.cat((class_token, x), dim=1)

        # 16. Add position embedding to patch embedding (equation 1)
        x = self.position_eembedding + x

        # 17. Run embedding dropout (Appendix B.1)
        x = self.embedding_dropout(x)

        # 18. Pass patch, position and class embedding through transformer encoder layers (equations 2 & 3)
        x = self.transformer_encoder(x)

        # 19. Put 0 index logit through classifier (equation 4)
        x = self.classifer_layer(x[:, 0]) # run on each sample in a batch at 0 index

        return x