File size: 5,184 Bytes
ff5b2f1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 | 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
|