File size: 4,052 Bytes
ff0fadf | 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 | import torch
import torch.nn as nn
import torch.nn.functional as F
from timm.layers import trunc_normal_
from onescience.modules.mlp.MLP import StandardMLP
from onescience.modules.transformer.gnot_transformer_block import GNOTTransformerBlock
from onescience.modules.embedding import timestep_embedding, unified_pos_embedding
class Model(nn.Module):
"""
GNOT (General Neural Operator Transformer) 模型。
"""
def __init__(self, args, device, n_experts=3):
super(Model, self).__init__()
self.__name__ = "GNOT"
self.args = args
# 1. Embedding & Preprocessing
# -----------------------------------------------------------
if args.unified_pos and args.geotype != "unstructured":
self.pos = unified_pos_embedding(args.shapelist, args.ref, device=device)
dim_x = args.ref ** len(args.shapelist)
dim_z = args.fun_dim + args.ref ** len(args.shapelist)
else:
dim_x = args.space_dim
dim_z = args.fun_dim + args.space_dim
self.preprocess_x = StandardMLP(
input_dim=dim_x,
output_dim=args.n_hidden,
hidden_dims=[args.n_hidden * 2],
activation=args.act,
use_bias=True
)
self.preprocess_z = StandardMLP(
input_dim=dim_z,
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 (MoE Style)
# -----------------------------------------------------------
self.blocks = nn.ModuleList([
GNOTTransformerBlock(
num_heads=args.n_heads,
hidden_dim=args.n_hidden,
dropout=args.dropout,
act=args.act,
mlp_ratio=args.mlp_ratio,
space_dim=args.space_dim,
n_experts=n_experts,
)
for _ in range(args.n_layers)
])
self.placeholder = nn.Parameter(
(1 / (args.n_hidden)) * torch.rand(args.n_hidden, dtype=torch.float)
)
# 3. Projectors (Decoder)
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 forward(self, x, fx, T=None, geo=None):
pos = x
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_z(fx)
else:
fx = self.preprocess_z(x)
fx = fx + self.placeholder[None, None, :]
x = self.preprocess_x(x) # x here becomes embedding of geometric info
if T is not None:
Time_emb = timestep_embedding(T, self.args.n_hidden) # (B, C)
Time_emb = self.time_fc(Time_emb)
if Time_emb.ndim == 2:
Time_emb = Time_emb.unsqueeze(1) # (B, 1, C)
fx = fx + Time_emb
for block in self.blocks:
# GNOT block 需要三个参数: x(geo_emb), fx(phys_emb), pos(coords)
fx = block(x, fx, pos)
fx = self.fc1(fx)
fx = F.gelu(fx)
fx = self.fc2(fx)
return fx
|