OneScience-Group/FourCastNet

1import numpy as np2import torch3import torch.nn as nn4from timm.models.layers import trunc_normal_5from einops import rearrange6 7from onescience.modules.embedding.fourcastnetembedding import FourCastNetEmbedding8from onescience.modules.fuser.fourcastnetfuser import FourCastNetFuser9 10 11class FourCastNet(nn.Module):12    """13    FourCastNet 的主模型实现。14 15    该模型使用以下组件完成输入编码与主干特征提取：16 17    - `OneEmbedding(style="FourCastNetEmbedding")`18      - 将二维气象场切分为二维 patch token 序列19    - `OneFuser(style="FourCastNetFuser")`20      - 在二维 patch 网格上重复执行 AFNO 频域混合与 MLP 通道混合21 22    在当前实现中：23 24    - 输入为二维单时刻气象场 `(Batch, Channels, Height, Width)`25    - patch embedding 输出会加上可学习位置编码26    - token 序列随后恢复成 `(PatchGridHeight, PatchGridWidth)` 二维网格27    - 多层 `FourCastNetFuser` 在 patch 网格上完成主干特征提取28    - 最终通过线性头恢复回目标变量场29 30    Args:31        img_size (tuple[int, int]):32            输入空间尺寸 `(Height, Width)`。33        patch_size (tuple[int, int]):34            patch 切分尺寸 `(PatchHeight, PatchWidth)`。35        in_chans (int):36            输入变量通道数。37        out_chans (int):38            输出变量通道数。39        embed_dim (int):40            patch embedding 特征维度。41        depth (int):42            主干 `FourCastNetFuser` 堆叠层数。43        mlp_ratio (float):44            每层 MLP 隐层放大倍数。45        drop_rate (float):46            dropout 比例。47        drop_path_rate (float):48            按层递增的 Stochastic Depth 最大比例。49        num_blocks (int):50            AFNO 的通道分块数。51        sparsity_threshold (float):52            AFNO 的 soft shrink 阈值。53        hard_thresholding_fraction (float):54            AFNO 保留的频率模式比例。55    """56 57    def __init__(58        self,59        img_size=(720, 1440),60        patch_size=(8, 8),61        in_chans=19,62        out_chans=19,63        embed_dim=768,64        depth=12,65        mlp_ratio=4.0,66        drop_rate=0.0,67        drop_path_rate=0.0,68        num_blocks=8,69        sparsity_threshold=0.01,70        hard_thresholding_fraction=1.0,71    ):72        super().__init__()73        self.img_size = img_size74        self.patch_size = patch_size75        self.in_chans = in_chans76        self.out_chans = out_chans77        self.num_features = self.embed_dim = embed_dim78        self.num_blocks = num_blocks79 80        num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0])81        drop_path = np.linspace(0, drop_path_rate, depth).tolist()82 83        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim))84        self.pos_drop = nn.Dropout(p=drop_rate)85        self.patch_grid_height = img_size[0] // self.patch_size[0]86        self.patch_grid_width = img_size[1] // self.patch_size[1]87 88        self.patch_embed = FourCastNetEmbedding(89            img_size=img_size,90            patch_size=patch_size,91            in_chans=in_chans,92            embed_dim=embed_dim,93        )94 95        self.blocks = nn.ModuleList([96            FourCastNetFuser(97                dim=embed_dim,98                mlp_ratio=mlp_ratio,99                drop=drop_rate,100                drop_path=drop_path[i],101                num_blocks=num_blocks,102                sparsity_threshold=sparsity_threshold,103                hard_thresholding_fraction=hard_thresholding_fraction,104            )105            for i in range(depth)106        ])107 108        self.head = nn.Linear(109            embed_dim,110            self.out_chans * self.patch_size[0] * self.patch_size[1],111            bias=False,112        )113 114        trunc_normal_(self.pos_embed, std=0.02)115        self.apply(self._init_weights)116 117    def _init_weights(self, m):118        if isinstance(m, nn.Linear):119            trunc_normal_(m.weight, std=0.02)120            if isinstance(m, nn.Linear) and m.bias is not None:121                nn.init.constant_(m.bias, 0)122        elif isinstance(m, nn.LayerNorm):123            nn.init.constant_(m.bias, 0)124            nn.init.constant_(m.weight, 1.0)125 126    @torch.jit.ignore127    def no_weight_decay(self):128        return {'pos_embed', 'cls_token'}129 130 131    def forward(self, x):132        """133        Args:134            x (torch.Tensor):135                输入张量，形状为 `(Batch, Channels, Height, Width)`。136 137        Returns:138            torch.Tensor:139                输出张量，形状为 `(Batch, out_chans, Height, Width)`。140        """141        Batch = x.shape[0]142 143        x = self.patch_embed(x)144        x = x + self.pos_embed145        x = self.pos_drop(x)146 147        x = x.reshape(Batch, self.patch_grid_height, self.patch_grid_width, self.embed_dim)148        for blk in self.blocks:149            x = blk(x)150 151        x = self.head(x)152        x = rearrange(153            x,154            "b h w (p1 p2 c_out) -> b c_out (h p1) (w p2)",155            p1=self.patch_size[0],156            p2=self.patch_size[1],157            h=self.patch_grid_height,158            w=self.patch_grid_width,159        )160        return x161