model.py

import torch
import torch.nn as nn
import torch.nn.functional as F
from einops.layers.torch import Rearrange

class AttentionPool(nn.Module):
    def __init__(self, dim, pool_size=2):
        super().__init__()
        self.pool_size = pool_size
        # (n p ) are length of sequence
        self.pool_fn = Rearrange("b d (n p) -> b d n p", p=pool_size)
        self.to_attn_logits = nn.Conv2d(dim, dim, 1, bias=False)

    def forward(self, x):
        b, _, n = x.shape
        remainder = n % self.pool_size
        needs_padding = remainder > 0
        if needs_padding:
            x = F.pad(x, (0, remainder), value=0)
            mask = torch.zeros((b, 1, n), dtype=torch.bool, device=x.device)
            mask = F.pad(mask, (0, remainder), value=True)
        x = self.pool_fn(x)
        logits = self.to_attn_logits(x)
        if needs_padding:
            mask_value = -torch.finfo(logits.dtype).max
            logits = logits.masked_fill(self.pool_fn(mask), mask_value)
        attn = logits.softmax(dim=-1)
        return (x * attn).sum(dim=-1)

class ConvBlock(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding):
        super(ConvBlock, self).__init__()
        self.conv = nn.Conv1d(in_channels=in_channels, 
                              out_channels=out_channels, 
                              kernel_size=kernel_size, 
                              stride=stride, 
                              padding=padding)
        self.batch_norm = nn.SyncBatchNorm(num_features=out_channels)
        self.relu = nn.ReLU()
    
    """
    Now we'll build the forward function which defines the path to input tensor
    meaning that we tell the tensor the sequence of layers you're going through 
    Takecare the name of forward function is sensitive so you have to name forward not any thing else
    """
    def forward(self, x):
        x = self.conv(x)
        x = self.batch_norm(x)
        x = self.relu(x)
        
        return x
    
class Inception1DBlock(nn.Module):
    def __init__(self, in_channels, filter_1x1, filter_3x3, filter_5x5, filter_3x3_Reduce,filter_5x5_Reduce, pooling):
        super(Inception1DBlock, self).__init__()
        # The In_Channels are the depth of tensor coming from previous layer
        # First block contains only filters with kernel size 1x1
        self.block_1 = nn.Sequential(
            ConvBlock(in_channels, filter_1x1, kernel_size=1, stride=1, padding=0),
            AttentionPool(filter_1x1, pool_size=2), # Attention Pooling layer
            )
        
        # Second Block contains filters with kernel size 1x1 followed by 3x3
        self.block_2 = nn.Sequential(
            ConvBlock(in_channels, filter_3x3_Reduce, kernel_size=1, stride=1, padding=0),
            ConvBlock(filter_3x3_Reduce, filter_3x3, kernel_size=3, stride=1, padding=1),
            AttentionPool(filter_3x3, pool_size=2), # Attention Pooling layer
        )
        
        # Third Block same as second block unless we'll replace the 3x3 filter with 5x5 
        self.block_3 = nn.Sequential(
            ConvBlock(in_channels, filter_5x5_Reduce, kernel_size=1, stride=1, padding=0),
            ConvBlock(filter_5x5_Reduce, filter_5x5, kernel_size=5, stride=1, padding=2),
            AttentionPool(filter_5x5, pool_size=2), # Attention Pooling layer
        )
        
        # Fourth Block contains maxpooling layer followed by 1x1 filter
        self.block_4 = nn.Sequential(
            nn.MaxPool1d(kernel_size=3, stride=1, padding=1),
            ConvBlock(in_channels, pooling, kernel_size=1, stride=1, padding=0),
            AttentionPool(pooling, pool_size=2), # Attention Pooling layer
        )

    def forward(self, x):
        out_1 = self.block_1(x)
        out_2 = self.block_2(x)
        out_3 = self.block_3(x)
        out_4 = self.block_4(x)
        
        conc_outs = torch.cat([out_1, out_2, out_3, out_4], dim=1) #dim=1 because we want to concatenate in the depth dimension
        return conc_outs
    

class InceptionNet1D(nn.Module):
    def __init__(self, num_classes, in_channels=4):
        super(InceptionNet1D, self).__init__()
        self.conv1 = ConvBlock(in_channels=in_channels, out_channels=64, kernel_size=7, stride=2, padding=3)
        # self.maxpool1 = nn.MaxPool1d(kernel_size=3, stride=2, padding=0, ceil_mode=True)
        self.conv2 = ConvBlock(in_channels=64, out_channels=64, kernel_size=1, stride=1, padding=0)
        self.conv3 = ConvBlock(in_channels=64, out_channels=192, kernel_size=3, stride=1, padding=1)
        # self.maxpool2 = nn.MaxPool1d(kernel_size=3, stride=2, padding=0, ceil_mode=True)
        self.inception_3a = Inception1DBlock(in_channels=192, filter_1x1=64, filter_3x3=128
                                          , filter_5x5=32, filter_3x3_Reduce=96, 
                                           filter_5x5_Reduce=16, pooling=32)
        
        self.inception_3b = Inception1DBlock(in_channels=256, filter_1x1=128, filter_3x3=192
                                          , filter_5x5=96, filter_3x3_Reduce=128, 
                                           filter_5x5_Reduce=32, pooling=64)
        
        # self.maxpool3 = nn.MaxPool1d(kernel_size=3, stride=2, padding=0, ceil_mode=True)
        self.inception_4a = Inception1DBlock(in_channels=480, filter_1x1=192, filter_3x3=208
                                          , filter_5x5=48, filter_3x3_Reduce=96, 
                                           filter_5x5_Reduce=16, pooling=64)
        
        self.inception_4b = Inception1DBlock(in_channels=512, filter_1x1=160, filter_3x3=224
                                          , filter_5x5=64, filter_3x3_Reduce=112, 
                                           filter_5x5_Reduce=24, pooling=64)
        
        
        self.inception_4c = Inception1DBlock(in_channels=512, filter_1x1=128, filter_3x3=256
                                          , filter_5x5=64, filter_3x3_Reduce=128, 
                                           filter_5x5_Reduce=24, pooling=64)
        
       
        self.inception_4d = Inception1DBlock(in_channels=512, filter_1x1=112, filter_3x3=288
                                          , filter_5x5=64, filter_3x3_Reduce=144, 
                                           filter_5x5_Reduce=32, pooling=64) 
        
        self.inception_4e = Inception1DBlock(in_channels=528, filter_1x1=256, filter_3x3=320
                                          , filter_5x5=128, filter_3x3_Reduce=160, 
                                           filter_5x5_Reduce=32, pooling=128) 
        
        # self.maxpool4 = nn.MaxPool1d(kernel_size=3, stride=2, padding=0, ceil_mode=True)
        self.inception_5a = Inception1DBlock(in_channels=832, filter_1x1=256, filter_3x3=320
                                          , filter_5x5=128, filter_3x3_Reduce=160, 
                                           filter_5x5_Reduce=32, pooling=128) 
        
        self.inception_5b = Inception1DBlock(in_channels=832, filter_1x1=384, filter_3x3=384
                                          , filter_5x5=128, filter_3x3_Reduce=192, 
                                           filter_5x5_Reduce=48, pooling=128) 
        
        self.avgpool = nn.AdaptiveAvgPool1d(output_size=1)
        self.dropout = nn.Dropout(p=0.4)
        self.fc = nn.Linear(in_features=1024, out_features=num_classes)
        
        
    def forward(self, x):
        out = self.conv1(x)
        # out = self.maxpool1(out)
        out = self.conv2(out)
        out = self.conv3(out)
        # out = self.maxpool2(out)

        out = self.inception_3a(out)
        out = self.inception_3b(out)
        # out = self.maxpool3(out)

        out = self.inception_4a(out)
        out = self.inception_4b(out)
        out = self.inception_4c(out)
        out = self.inception_4d(out)
        out = self.inception_4e(out)
        # out = self.maxpool4(out)

        out = self.inception_5a(out)
        out = self.inception_5b(out)
        
        out = self.avgpool(out)
        out = torch.flatten(out, 1)
        out = self.dropout(out)
        out = self.fc(out)
        
        return out