X_SSL_Net/lib/modules/attentions_2d.py

"""
Circulant Attention 2D.

核心思想: 自注意力矩阵近似 BC CB 结构，通过 2D FFT 在 O(N log N) 时间内计算。
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Literal

try:
    import ptwt
except ImportError as exc:
    raise ImportError(
        "wavelet_fft requires ptwt. Install it before importing this package."
    ) from exc

from .layers_2d import Scale


class ComplexLinear(nn.Linear):
    def __init__(self, in_features, out_features, device=None, dtype=None):
        super().__init__(
            in_features, out_features, bias=False, device=device, dtype=dtype
        )

    def forward(self, inp):
        x = torch.view_as_real(inp).transpose(-2, -1)
        x = F.linear(x, self.weight).transpose(-2, -1)
        if x.dtype != torch.float32:
            x = x.to(torch.float32)
        return torch.view_as_complex(x.contiguous())


class CirculantAttention2d(nn.Module):
    def __init__(self, dim, proj_drop=0.0):
        super().__init__()
        self.qkv = ComplexLinear(dim, dim * 3)
        self.gate = nn.Sequential(nn.Linear(dim, dim), nn.SiLU())
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)

    def forward(self, x):
        b, c, h, w = x.shape
        spatial_perm = [0, 2, 3, 1]
        spatial = x.permute(spatial_perm).contiguous()
        gate = self.gate(spatial.reshape(b, h * w, c)).reshape(b, h, w, c)
        freq = torch.fft.rfft2(spatial, dim=(1, 2), norm="ortho")
        qkv = self.qkv(freq)
        q, k, v = torch.chunk(qkv, chunks=3, dim=-1)
        attn = torch.conj(q) * k
        attn = torch.fft.irfft2(attn, s=(h, w), dim=(1, 2), norm="ortho")
        attn = attn.reshape(b, h * w, c).softmax(dim=1).reshape(b, h, w, c)
        attn = torch.fft.rfft2(attn, dim=(1, 2))
        out = torch.conj(attn) * v
        out = torch.fft.irfft2(out, s=(h, w), dim=(1, 2), norm="ortho")
        out = out.reshape(b, h * w, c) * gate.reshape(b, h * w, c)
        out = self.proj_drop(self.proj(out))
        return out.transpose(1, 2).reshape(b, c, h, w)


class WaveletAttentionGlobalBranch2d(nn.Module):
    def __init__(
        self,
        in_channels,
        kernel_size=5,
        stride=1,
        wt_levels=1,
        wt_type="db1",
        wt_mode: Literal[
            "constant", "zero", "reflect", "periodic", "symmetric"
        ] = "zero",
        proj_drop=0.0,
    ):
        super().__init__()
        if in_channels <= 0:
            raise ValueError("in_channels must be positive.")

        self.in_channels = in_channels
        self.wt_levels = wt_levels
        self.stride = stride

        self.wavelet = wt_type
        self.wt_mode = wt_mode

        self.global_attn = CirculantAttention2d(in_channels, proj_drop=proj_drop)
        self.base_scale = Scale((1, in_channels, 1, 1))

        self.wavelet_convs = nn.ModuleList(
            [
                nn.Conv2d(
                    in_channels * 4,
                    in_channels * 4,
                    kernel_size,
                    1,
                    kernel_size // 2,
                    groups=in_channels * 4,
                    bias=False,
                )
                for _ in range(wt_levels)
            ]
        )
        self.wavelet_scale = nn.ModuleList(
            [
                Scale((1, in_channels * 4, 1, 1), init_scale=0.1)
                for _ in range(wt_levels)
            ]
        )

        if stride > 1:
            self.register_buffer(
                "stride_filter", torch.ones(in_channels, 1, 1, 1), persistent=False
            )
        else:
            self.stride_filter = None

    def forward(self, x):
        low_levels, high_levels, shapes_in_levels = [], [], []
        curr_low = x

        for level in range(self.wt_levels):
            shapes_in_levels.append(curr_low.shape[-2:])
            coeffs = ptwt.wavedec2(curr_low, self.wavelet, mode=self.wt_mode, level=1)
            low = coeffs[0]
            detail = coeffs[1]
            high = torch.stack(
                [detail.horizontal, detail.vertical, detail.diagonal], dim=2
            )
            bands = torch.cat([low.unsqueeze(2), high], dim=2)
            b, c, _, h_half, w_half = bands.shape
            bands = bands.reshape(b, c * 4, h_half, w_half)
            bands = self.wavelet_scale[level](self.wavelet_convs[level](bands))
            bands = bands.reshape(b, c, 4, h_half, w_half)
            low_levels.append(bands[:, :, 0, :, :])
            high_levels.append(bands[:, :, 1:4, :, :])
            curr_low = low

        wavelet_out = x
        if self.wt_levels > 0:
            next_low = None
            for level in range(self.wt_levels - 1, -1, -1):
                low = low_levels.pop()
                high = high_levels.pop()
                height, width = shapes_in_levels.pop()
                if next_low is not None:
                    low = low + next_low
                cH, cV, cD = high.unbind(dim=2)
                next_low = ptwt.waverec2(
                    (low, ptwt.constants.WaveletDetailTuple2d(cH, cV, cD)), self.wavelet
                )
                next_low = next_low[:, :, :height, :width]
            wavelet_out = next_low

        out = self.base_scale(self.global_attn(x)) + wavelet_out
        if self.stride_filter is not None:
            out = F.conv2d(
                out, self.stride_filter, stride=self.stride, groups=self.in_channels
            )
        return out