| import numpy as np |
| import torch |
| import torch.nn as nn |
| from timm.models.layers import trunc_normal_ |
| from einops import rearrange |
|
|
| from onescience.modules.embedding.fourcastnetembedding import FourCastNetEmbedding |
| from onescience.modules.fuser.fourcastnetfuser import FourCastNetFuser |
|
|
|
|
| class FourCastNet(nn.Module): |
| """ |
| FourCastNet 的主模型实现。 |
| |
| 该模型使用以下组件完成输入编码与主干特征提取： |
| |
| - `OneEmbedding(style="FourCastNetEmbedding")` |
| - 将二维气象场切分为二维 patch token 序列 |
| - `OneFuser(style="FourCastNetFuser")` |
| - 在二维 patch 网格上重复执行 AFNO 频域混合与 MLP 通道混合 |
| |
| 在当前实现中： |
| |
| - 输入为二维单时刻气象场 `(Batch, Channels, Height, Width)` |
| - patch embedding 输出会加上可学习位置编码 |
| - token 序列随后恢复成 `(PatchGridHeight, PatchGridWidth)` 二维网格 |
| - 多层 `FourCastNetFuser` 在 patch 网格上完成主干特征提取 |
| - 最终通过线性头恢复回目标变量场 |
| |
| Args: |
| img_size (tuple[int, int]): |
| 输入空间尺寸 `(Height, Width)`。 |
| patch_size (tuple[int, int]): |
| patch 切分尺寸 `(PatchHeight, PatchWidth)`。 |
| in_chans (int): |
| 输入变量通道数。 |
| out_chans (int): |
| 输出变量通道数。 |
| embed_dim (int): |
| patch embedding 特征维度。 |
| depth (int): |
| 主干 `FourCastNetFuser` 堆叠层数。 |
| mlp_ratio (float): |
| 每层 MLP 隐层放大倍数。 |
| drop_rate (float): |
| dropout 比例。 |
| drop_path_rate (float): |
| 按层递增的 Stochastic Depth 最大比例。 |
| num_blocks (int): |
| AFNO 的通道分块数。 |
| sparsity_threshold (float): |
| AFNO 的 soft shrink 阈值。 |
| hard_thresholding_fraction (float): |
| AFNO 保留的频率模式比例。 |
| """ |
|
|
| def __init__( |
| self, |
| img_size=(720, 1440), |
| patch_size=(8, 8), |
| in_chans=19, |
| out_chans=19, |
| embed_dim=768, |
| depth=12, |
| mlp_ratio=4.0, |
| drop_rate=0.0, |
| drop_path_rate=0.0, |
| num_blocks=8, |
| sparsity_threshold=0.01, |
| hard_thresholding_fraction=1.0, |
| ): |
| super().__init__() |
| self.img_size = img_size |
| self.patch_size = patch_size |
| self.in_chans = in_chans |
| self.out_chans = out_chans |
| self.num_features = self.embed_dim = embed_dim |
| self.num_blocks = num_blocks |
|
|
| num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0]) |
| drop_path = np.linspace(0, drop_path_rate, depth).tolist() |
|
|
| self.pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim)) |
| self.pos_drop = nn.Dropout(p=drop_rate) |
| self.patch_grid_height = img_size[0] // self.patch_size[0] |
| self.patch_grid_width = img_size[1] // self.patch_size[1] |
|
|
| self.patch_embed = FourCastNetEmbedding( |
| img_size=img_size, |
| patch_size=patch_size, |
| in_chans=in_chans, |
| embed_dim=embed_dim, |
| ) |
|
|
| self.blocks = nn.ModuleList([ |
| FourCastNetFuser( |
| dim=embed_dim, |
| mlp_ratio=mlp_ratio, |
| drop=drop_rate, |
| drop_path=drop_path[i], |
| num_blocks=num_blocks, |
| sparsity_threshold=sparsity_threshold, |
| hard_thresholding_fraction=hard_thresholding_fraction, |
| ) |
| for i in range(depth) |
| ]) |
|
|
| self.head = nn.Linear( |
| embed_dim, |
| self.out_chans * self.patch_size[0] * self.patch_size[1], |
| bias=False, |
| ) |
|
|
| trunc_normal_(self.pos_embed, std=0.02) |
| 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.init.constant_(m.bias, 0) |
| nn.init.constant_(m.weight, 1.0) |
|
|
| @torch.jit.ignore |
| def no_weight_decay(self): |
| return {'pos_embed', 'cls_token'} |
|
|
|
|
| def forward(self, x): |
| """ |
| Args: |
| x (torch.Tensor): |
| 输入张量，形状为 `(Batch, Channels, Height, Width)`。 |
| |
| Returns: |
| torch.Tensor: |
| 输出张量，形状为 `(Batch, out_chans, Height, Width)`。 |
| """ |
| Batch = x.shape[0] |
|
|
| x = self.patch_embed(x) |
| x = x + self.pos_embed |
| x = self.pos_drop(x) |
|
|
| x = x.reshape(Batch, self.patch_grid_height, self.patch_grid_width, self.embed_dim) |
| for blk in self.blocks: |
| x = blk(x) |
|
|
| x = self.head(x) |
| x = rearrange( |
| x, |
| "b h w (p1 p2 c_out) -> b c_out (h p1) (w p2)", |
| p1=self.patch_size[0], |
| p2=self.patch_size[1], |
| h=self.patch_grid_height, |
| w=self.patch_grid_width, |
| ) |
| return x |
|
|
