model/Swin_Transformer.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.checkpoint as checkpoint
from timm.layers import trunc_normal_
from onescience.modules.mlp.MLP import StandardMLP
from onescience.modules.transformer.SwinTransformerBlock import SwinTransformerBlock
from onescience.modules.embedding import timestep_embedding, unified_pos_embedding

class BasicLayer(nn.Module):
    """
    Swin Transformer Layer (Stage).
    封装了多个 SwinTransformerBlock。
    """
    def __init__(
        self,
        dim,
        input_resolution,
        depth,
        num_heads,
        window_size,
        mlp_ratio=4.0,
        qkv_bias=True,
        qk_scale=None,
        drop=0.0,
        attn_drop=0.0,
        drop_path=0.0,
        norm_layer=nn.LayerNorm,
        downsample=None,
        use_checkpoint=False,
        fused_window_process=False,
    ):
        super().__init__()
        self.dim = dim
        self.input_resolution = input_resolution
        self.depth = depth
        self.use_checkpoint = use_checkpoint

        # build blocks
        self.blocks = nn.ModuleList([
            SwinTransformerBlock(
                dim=dim,
                input_resolution=input_resolution,
                num_heads=num_heads,
                window_size=window_size,
                shift_size=0 if (i % 2 == 0) else window_size // 2,
                mlp_ratio=mlp_ratio,
                qkv_bias=qkv_bias,
                qk_scale=qk_scale,
                drop=drop,
                attn_drop=attn_drop,
                drop_path=(drop_path[i] if isinstance(drop_path, list) else drop_path),
                norm_layer=norm_layer,
                fused_window_process=fused_window_process,
            )
            for i in range(depth)
        ])

        # patch merging layer
        if downsample is not None:
            self.downsample = downsample(
                input_resolution, dim=dim, norm_layer=norm_layer
            )
        else:
            self.downsample = None

    def forward(self, x):
        for blk in self.blocks:
            if self.use_checkpoint:
                x = checkpoint.checkpoint(blk, x)
            else:
                x = blk(x)
        if self.downsample is not None:
            x = self.downsample(x)
        return x


class Model(nn.Module):
    """
    Swin Transformer 主模型。
    """
    def __init__(self, args, device, window_size=4):
        super(Model, self).__init__()
        self.__name__ = "SwinTransformer"
        self.args = args
        
        if args.geotype != "structured_2D":
            raise ValueError("Swin Transformer only supports Structured 2D geometry")

        # 1. Embedding & Preprocessing
        if args.unified_pos: 
            self.pos = unified_pos_embedding(args.shapelist, args.ref, device=device)
            in_dim = args.fun_dim + args.ref ** len(args.shapelist)
        else:
            in_dim = args.fun_dim + args.space_dim

        self.preprocess = StandardMLP(
            input_dim=in_dim,
            output_dim=args.n_hidden,
            hidden_dims=[args.n_hidden * 2],
            activation=args.act,
            use_bias=True
        )
        
        if args.time_input:
            self.time_fc = nn.Sequential(
                nn.Linear(args.n_hidden, args.n_hidden),
                nn.SiLU(),
                nn.Linear(args.n_hidden, args.n_hidden),
            )
            
        self.placeholder = nn.Parameter(
            (1 / (args.n_hidden)) * torch.rand(args.n_hidden, dtype=torch.float)
        )
        
        self.padding = [
            (window_size - size % window_size) % window_size for size in args.shapelist
        ]
        self.augmented_resolution = [
            (self.padding[i] + args.shapelist[i]) for i in range(len(self.padding))
        ]
        
        # 3. Swin Layers
        self.blocks = nn.ModuleList([
            BasicLayer(
                dim=args.n_hidden,
                input_resolution=self.augmented_resolution,
                depth=2,
                num_heads=args.n_heads,
                window_size=window_size,
            )
            for _ in range(args.n_layers)
        ])
        
        # 4. Projectors
        self.fc1 = nn.Linear(args.n_hidden, args.n_hidden * 2)
        self.fc2 = nn.Linear(args.n_hidden * 2, args.out_dim)
        
        self.initialize_weights()

    def initialize_weights(self):
        self.apply(self._init_weights)

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            trunc_normal_(m.weight, std=0.02)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, (nn.LayerNorm, nn.BatchNorm1d)):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0)

    def structured_geo(self, x, fx, T=None):
        B, N, _ = x.shape
        if self.args.unified_pos:
            x = self.pos.repeat(x.shape[0], 1, 1)
            
        if fx is not None:
            fx = torch.cat((x, fx), -1)
            fx = self.preprocess(fx)
        else:
            fx = self.preprocess(x)
            
        fx = fx + self.placeholder[None, None, :]

        if T is not None:
            Time_emb = timestep_embedding(T, self.args.n_hidden)
            Time_emb = self.time_fc(Time_emb)
            if Time_emb.ndim == 2:
                Time_emb = Time_emb.unsqueeze(1)
            fx = fx + Time_emb
            
        ## aug shape
        fx = fx.permute(0, 2, 1).reshape(B, self.args.n_hidden, *self.args.shapelist)
        
        if not all(item == 0 for item in self.padding):
            if len(self.args.shapelist) == 2:
                fx = F.pad(fx, [0, self.padding[1], 0, self.padding[0]])
            elif len(self.args.shapelist) == 3:
                fx = F.pad(fx, [0, self.padding[2], 0, self.padding[1], 0, self.padding[0]])
                
        fx = fx.reshape(B, self.args.n_hidden, -1).permute(0, 2, 1)
        
        ## swin transformer
        for block in self.blocks:
            fx = block(fx)
            
        ## back to original shape
        fx = fx.permute(0, 2, 1).reshape(
            B, self.args.n_hidden, *self.augmented_resolution
        )
        
        if not all(item == 0 for item in self.padding):
            if len(self.args.shapelist) == 2:
                fx = fx[..., : -self.padding[0], : -self.padding[1]]
            elif len(self.args.shapelist) == 3:
                fx = fx[..., : -self.padding[0], : -self.padding[1], : -self.padding[2]]
                
        fx = fx.reshape(B, self.args.n_hidden, -1).permute(0, 2, 1)
        
        ## projection
        fx = self.fc1(fx)
        fx = F.gelu(fx)
        fx = self.fc2(fx)
        return fx

    def forward(self, x, fx, T=None, geo=None):
        if self.args.geotype == "structured_2D":
            return self.structured_geo(x, fx, T)
        else:
            raise ValueError("Swin Transformer only supports Structured 2D geometry")