mecnn.py

import torch
import torch.nn as nn

class MECNN(nn.Module):
    """
    Multiscale Entropy Convolutional Neural Network (MECNN)
    Architecture detailed in Section III.F and Table 2.
    """
    def __init__(self, num_channels=22, num_scales=24, num_classes=4, dropout_rate=0.5):
        super(MECNN, self).__init__()
        
        # 1. Scale-wise Convolution (Kernel: 1x3, 32 filters)
        # Extracts local patterns along the temporal scale axis independently.
        self.conv_scale = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=(1, 3), stride=(1, 1))
        self.bn_scale = nn.BatchNorm2d(32)
        self.relu_scale = nn.ReLU()
        
        # 2. Spatial Convolution (Kernel: Cx1, 64 filters)
        # Collapses spatial dimension to capture spatial correlations across electrodes.
        self.conv_spatial = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=(num_channels, 1), stride=(1, 1))
        self.bn_spatial = nn.BatchNorm2d(64)
        self.relu_spatial = nn.ReLU()
        
        # 3. Average Pooling (Kernel: 1x2)
        self.pool = nn.AvgPool2d(kernel_size=(1, 2), stride=(1, 2))
        
        # Calculate flattened dimension dynamically based on input scales
        # After conv_scale: Width = S - 2
        # After avg_pool: Width = (S - 2) // 2
        pool_out_w = (num_scales - 2) // 2
        self.flattened_dim = 64 * 1 * pool_out_w # e.g., 64 * 1 * 11 = 704
        
        # 4. Fully Connected Layers
        self.fc1 = nn.Linear(self.flattened_dim, 128)
        self.relu_fc1 = nn.ReLU()
        self.dropout = nn.Dropout(p=dropout_rate)
        
        # Output layer
        self.fc_out = nn.Linear(128, num_classes)

    def forward(self, x):
        # Input shape: (Batch, 1, Channels, Scales)
        x = self.relu_scale(self.bn_scale(self.conv_scale(x)))
        x = self.relu_spatial(self.bn_spatial(self.conv_spatial(x)))
        x = self.pool(x)
        
        x = torch.flatten(x, start_dim=1)
        
        x = self.dropout(self.relu_fc1(self.fc1(x)))
        logits = self.fc_out(x)
        
        return logits

if __name__ == "__main__":
    # Sanity Check for Table 2 (Parameter count)
    model = MECNN(num_channels=22, num_scales=24, num_classes=4)
    dummy_input = torch.randn(2, 1, 22, 24)
    print(f"Output Shape: {model(dummy_input).shape}")
    
    # In PyTorch, BatchNorm running mean/var are not counted in requires_grad=True, 
    # but they are part of the total parameters (136,388) shown in Table 2.
    total_params = sum(p.numel() for p in model.parameters())
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f"Trainable Params: {trainable_params}, Total Params (incl. BN stats): {total_params}")