| # import torch | |
| # import torch.nn as nn | |
| # import torch.nn.functional as F | |
| # import numpy as np | |
| # from timm.layers import trunc_normal_ | |
| # from onescience.modules.layer.layers.Basic import MLP | |
| # from onescience.modules.attention.linearattention import LinearAttention | |
| # from onescience.modules.embedding import timestep_embedding, unified_pos_embedding | |
| # from einops import rearrange, repeat | |
| # from einops.layers.torch import Rearrange | |
| # class Galerkin_Transformer_block(nn.Module): | |
| # """Transformer encoder block.""" | |
| # def __init__( | |
| # self, | |
| # num_heads: int, | |
| # hidden_dim: int, | |
| # dropout: float, | |
| # act="gelu", | |
| # mlp_ratio=4, | |
| # last_layer=False, | |
| # out_dim=1, | |
| # ): | |
| # super().__init__() | |
| # self.last_layer = last_layer | |
| # self.ln_1 = nn.LayerNorm(hidden_dim) | |
| # self.ln_1a = nn.LayerNorm(hidden_dim) | |
| # self.Attn = LinearAttention( | |
| # hidden_dim, | |
| # heads=num_heads, | |
| # dim_head=hidden_dim // num_heads, | |
| # dropout=dropout, | |
| # attn_type="galerkin", | |
| # ) | |
| # self.ln_2 = nn.LayerNorm(hidden_dim) | |
| # self.mlp = MLP( | |
| # hidden_dim, | |
| # hidden_dim * mlp_ratio, | |
| # hidden_dim, | |
| # n_layers=0, | |
| # res=False, | |
| # act=act, | |
| # ) | |
| # if self.last_layer: | |
| # self.ln_3 = nn.LayerNorm(hidden_dim) | |
| # self.mlp2 = nn.Linear(hidden_dim, out_dim) | |
| # def forward(self, fx): | |
| # fx = self.Attn(self.ln_1(fx), self.ln_1a(fx)) + fx | |
| # fx = self.mlp(self.ln_2(fx)) + fx | |
| # if self.last_layer: | |
| # return self.mlp2(self.ln_3(fx)) | |
| # else: | |
| # return fx | |
| # class Model(nn.Module): | |
| # ## Galerkin_Transformer | |
| # def __init__(self, args, device): | |
| # super(Model, self).__init__() | |
| # self.__name__ = "Galerkin_Transformer" | |
| # self.args = args | |
| # ## embedding | |
| # if ( | |
| # args.unified_pos and args.geotype != "unstructured" | |
| # ): # only for structured mesh | |
| # self.pos = unified_pos_embedding(args.shapelist, args.ref, device=device) | |
| # self.preprocess = MLP( | |
| # args.fun_dim + args.ref ** len(args.shapelist), | |
| # args.n_hidden * 2, | |
| # args.n_hidden, | |
| # n_layers=0, | |
| # res=False, | |
| # act=args.act, | |
| # ) | |
| # else: | |
| # self.preprocess = MLP( | |
| # args.fun_dim + args.space_dim, | |
| # args.n_hidden * 2, | |
| # args.n_hidden, | |
| # n_layers=0, | |
| # res=False, | |
| # act=args.act, | |
| # ) | |
| # 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), | |
| # ) | |
| # ## models | |
| # self.blocks = nn.ModuleList( | |
| # [ | |
| # Galerkin_Transformer_block( | |
| # num_heads=args.n_heads, | |
| # hidden_dim=args.n_hidden, | |
| # dropout=args.dropout, | |
| # act=args.act, | |
| # mlp_ratio=args.mlp_ratio, | |
| # out_dim=args.out_dim, | |
| # last_layer=(_ == args.n_layers - 1), | |
| # ) | |
| # for _ in range(args.n_layers) | |
| # ] | |
| # ) | |
| # self.placeholder = nn.Parameter( | |
| # (1 / (args.n_hidden)) * torch.rand(args.n_hidden, dtype=torch.float) | |
| # ) | |
| # 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 forward(self, x, fx, T=None, geo=None): | |
| # 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) # (B, C) | |
| # Time_emb = self.time_fc(Time_emb) # (B, C) | |
| # if Time_emb.ndim == 2: | |
| # Time_emb = Time_emb.unsqueeze(1) # (B, 1, C) | |
| # fx = fx + Time_emb # Broadcasting: (B, N, C) + (B, 1, C) -> (B, N, C) | |
| # for block in self.blocks: | |
| # fx = block(fx) | |
| # return fx | |
| import torch | |
| import torch.nn as nn | |
| from timm.layers import trunc_normal_ | |
| # --- 引入模块工厂 --- | |
| from onescience.modules.mlp.MLP import StandardMLP | |
| from onescience.modules.transformer.galerkin_transformer_block import Galerkin_Transformer_block | |
| from onescience.modules.embedding import timestep_embedding, unified_pos_embedding | |
| class Model(nn.Module): | |
| """ | |
| Galerkin Transformer 模型。 | |
| 使用 Galerkin 线性注意力机制处理物理场数据的 Transformer 架构。 | |
| """ | |
| def __init__(self, args, device): | |
| super(Model, self).__init__() | |
| self.__name__ = "Galerkin_Transformer" | |
| self.args = args | |
| # 1. Embedding & Preprocessing | |
| # ----------------------------------------------------------- | |
| input_dim = args.fun_dim | |
| if args.unified_pos and args.geotype != "unstructured": | |
| self.pos = unified_pos_embedding(args.shapelist, args.ref, device=device) | |
| input_dim += args.ref ** len(args.shapelist) | |
| else: | |
| input_dim += args.space_dim | |
| # 对应原代码: MLP(input_dim, hidden*2, hidden, n_layers=0) | |
| self.preprocess = StandardMLP( | |
| input_dim=input_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), | |
| ) | |
| # 2. Transformer Blocks | |
| # ----------------------------------------------------------- | |
| # 使用工厂实例化 Galerkin_Transformer_block | |
| self.blocks = nn.ModuleList([ | |
| Galerkin_Transformer_block( | |
| num_heads=args.n_heads, | |
| hidden_dim=args.n_hidden, | |
| dropout=args.dropout, | |
| act=args.act, | |
| mlp_ratio=args.mlp_ratio, | |
| out_dim=args.out_dim, | |
| last_layer=(_ == args.n_layers - 1) | |
| ) | |
| for _ in range(args.n_layers) | |
| ]) | |
| self.placeholder = nn.Parameter( | |
| (1 / (args.n_hidden)) * torch.rand(args.n_hidden, dtype=torch.float) | |
| ) | |
| 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 forward(self, x, fx, T=None, geo=None): | |
| 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) # (B, C) | |
| Time_emb = self.time_fc(Time_emb) # (B, C) | |
| if Time_emb.ndim == 2: | |
| Time_emb = Time_emb.unsqueeze(1) # (B, 1, C) | |
| fx = fx + Time_emb # Broadcasting: (B, N, C) + (B, 1, C) -> (B, N, C) | |
| for block in self.blocks: | |
| fx = block(fx) | |
| return fx | |