function.py

import torch.nn as nn, torch

"""Sample code"""

############################################################################################################################################
### SimpleGate
class SimpleGate(nn.Module):
    def forward(self, x):
        x1, x2 = x.chunk(2, dim=1)
        return x1 * x2
    
### LayerNorm2d
class LayerNorm2d(nn.Module):
    def __init__(self, channels, eps=1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(channels))
        self.bias = nn.Parameter(torch.zeros(channels))
        self.eps = eps

    def forward(self, x):
        mean = x.mean(dim=1, keepdim=True)
        var = (x - mean).pow(2).mean(dim=1, keepdim=True)

        x = (x - mean) / torch.sqrt(var + self.eps)

        return (
            self.weight[:, None, None] * x +
            self.bias[:, None, None]
        )

# Multi-Dconv Head Transposed Attention (MDTA)
class BiasFreeLayerNorm(nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(channels))

    def forward(self, x):
        variance = x.var(dim=1, keepdim=True, unbiased=False)
        return x / torch.sqrt(variance + 1e-5) * self.weight[:, None, None]


class WithBiasLayerNorm(nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(channels))
        self.bias = nn.Parameter(torch.zeros(channels))

    def forward(self, x):
        mean = x.mean(dim=1, keepdim=True)
        variance = x.var(dim=1, keepdim=True, unbiased=False)
        x = (x - mean) / torch.sqrt(variance + 1e-5)
        return x * self.weight[:, None, None] + self.bias[:, None, None]


class LayerNorm2d_att(nn.Module):
    def __init__(self, channels, layer_norm_type="WithBias"):
        super().__init__()

        if layer_norm_type == "BiasFree":
            self.norm = BiasFreeLayerNorm(channels)
        elif layer_norm_type == "WithBias":
            self.norm = WithBiasLayerNorm(channels)
        else:
            raise ValueError("layer_norm_type must be 'BiasFree' or 'WithBias'")

    def forward(self, x):
        return self.norm(x)