CNN with gMLP for Feature Extraction and Classification

import torch
import torch.nn as nn
import torch.nn.functional as F
from gmlp import SpatialGatingUnit ,gMLPBlock,gMLP

class CNN(nn.Module):
    def __init__(self):  # 初始化方法
        super(CNN, self).__init__()  # 初始化方法
        # 第一层卷积层，输入为 (1,10000,12)
        self.conv1 = nn.Sequential(  # input shape (1,10000,12)
            nn.Conv2d(  # 二维卷积层
                in_channels=1,  # input height   # 输入通道数
                out_channels=5,  # n_filters     # 输出通道数（卷积核数量）
                kernel_size=(200, 3),  # filter size  # 卷积核大小
                stride=(50, 1),  # filter movement/step   # 卷积核移动步长
                padding=1,  # 填充大小
            ),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, padding=1),  # 最大池化层
        )
        # 第二层卷积层，输入为 (5,99,7)
        self.conv2 = nn.Sequential(  # input shape (5,99,7)
            nn.Conv2d(5, 10, (20, 2), (4, 1), 1),  # output shape   # 输出通道数为10，卷积核大小为(20,2)，步长为(4,1)，填充为1
            nn.ReLU(),  # activation
            nn.MaxPool2d(kernel_size=2),  # 最大池化层  # output shape (10,10,4)
        )

        self.gmlp = gMLP(d_model=10, d_ffn=20, seq_len=40, num_layers=6)
              #  gmlp(data).shape

        self.flatten = nn.Flatten()

        self.linear = nn.Linear(in_features=400, out_features=64)
        self.relu = nn.ReLU(inplace=True)
        self.dropout = nn.Dropout(p=0.4)
        self.out = nn.Linear(64, 6)  # 全连接层，输出6个类别

    # 定义一个名为CNN的类，继承自PyTorch的nn.Module

    def forward(self, x):  # 前向传播方法
        x = self.conv1(x)  # 通过第一个卷积层序列
        x = self.conv2(x)  # 通过第二个卷积层序列
        # x = x.view(x.size(0), -1)      # 将张量展平
        # 四维转三维
        N, C, H, W = x.shape
        x = x.view(N, C, H * W).permute(0, 2, 1)
        x = self.gmlp(x)
        print(x.shape)
        x = self.flatten(x)#换顺序了 先展平 在输入线性层中
        x = self.linear(x)
        feature = x  # 提取特征
        x = self.relu(x)
        x = self.dropout(x)
        output = self.out(x)  # 通过全连接层得到输出
        return feature, output  # 返回特征和输出

#
# x=torch.randn(10, 1, 10000, 12)
# net = CNN()
#
# _, out = net(x)
# print(out.shape)

class SpatialGatingUnit(nn.Module):
    def __init__(self, dim):
        super(SpatialGatingUnit, self).__init__()
        self.fc = nn.Linear(dim, dim)

    def forward(self, x):
        return torch.sigmoid(self.fc(x))

class gMLPBlock(nn.Module):
    def __init__(self, d_model, d_ffn):
        super(gMLPBlock, self).__init__()
        self.fc1 = nn.Linear(d_model, d_ffn)
        self.sgu = SpatialGatingUnit(d_ffn)
        self.fc2 = nn.Linear(d_ffn, d_model)

    def forward(self, x):
        x = self.fc1(x)
        gate = self.sgu(x)
        x = self.fc2(x) * gate
        return x

class gMLP(nn.Module):
    def __init__(self, d_model, d_ffn, seq_len, num_layers):
        super(gMLP, self).__init__()
        self.blocks = nn.ModuleList([gMLPBlock(d_model, d_ffn) for _ in range(num_layers)])

    def forward(self, x):
        for block in self.blocks:
            x = block(x)
        return x

gMLP (gated Multilayer Perceptron) is a type of neural network architecture that replaces the traditional self-attention mechanism in transformer models with a spatial gating unit. The spatial gating unit is responsible for capturing global dependencies in the input sequence.

To implement gMLP in your CNN model, you can create a separate gMLP class that consists of multiple gMLP blocks. Each gMLP block contains a SpatialGatingUnit followed by a fully connected layer.

This implementation uses a SpatialGatingUnit to capture global dependencies in the input sequence, replacing traditional self-attention in Transformers. This allows the model to learn more effective representations for downstream tasks, such as classification. The CNN architecture is designed to extract features from the input data, and the gMLP is used to further refine these features, leading to improved performance on the classification task.