File size: 4,372 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 118 119 120 121 122 123 124 125 126 127 128 129 130 131 | import torch
import torch.nn as nn
import numpy as np
from timm.layers import trunc_normal_
from onescience.modules.mlp.MLP import StandardMLP
from onescience.modules.transformer.Transolver_block import Transolver_block
from onescience.modules.embedding import timestep_embedding, unified_pos_embedding
class Model(nn.Module):
"""
Transolver 模型。
通过物理启发的切片机制 (Slicing) 解决 PDE 和物理场预测问题。
"""
def __init__(self, args, device):
super(Model, self).__init__()
self.__name__ = "Transolver"
self.args = args
## embedding
if (
args.unified_pos and args.geotype != "unstructured"
):
self.pos = unified_pos_embedding(args.shapelist, args.ref, device=device)
self.preprocess = StandardMLP(
input_dim=args.fun_dim + args.ref ** len(args.shapelist),
output_dim=args.n_hidden,
hidden_dims=[args.n_hidden * 2],
activation=args.act,
use_bias=True
)
else:
self.preprocess = StandardMLP(
input_dim=args.fun_dim + args.space_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),
)
## models
self.blocks = nn.ModuleList(
[
Transolver_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,
slice_num=args.slice_num,
last_layer=(_ == args.n_layers - 1),
geotype=args.geotype,
shapelist=args.shapelist,
)
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 structured_geo(self, x, fx, T=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)
Time_emb = self.time_fc(Time_emb)
if Time_emb.ndim == 2:
Time_emb = Time_emb.unsqueeze(1)
fx = fx + Time_emb
for block in self.blocks:
fx = block(fx)
return fx
def unstructured_geo(self, x, fx, T=None):
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)
Time_emb = self.time_fc(Time_emb)
if Time_emb.ndim == 2:
Time_emb = Time_emb.unsqueeze(1)
fx = fx + Time_emb
for block in self.blocks:
fx = block(fx)
return fx
def forward(self, x, fx, T=None, geo=None):
if self.args.geotype == "unstructured":
return self.unstructured_geo(x, fx, T)
else:
return self.structured_geo(x, fx, T)
|