steady-rans-surrogates / code /ezflow_v3 /baselines /Transolver-main /Car-Design-ShapeNetCar /models /Transolver.py
| import torch | |
| import numpy as np | |
| import torch.nn as nn | |
| from timm.models.layers import trunc_normal_ | |
| from einops import rearrange, repeat | |
| ACTIVATION = {'gelu': nn.GELU, 'tanh': nn.Tanh, 'sigmoid': nn.Sigmoid, 'relu': nn.ReLU, 'leaky_relu': nn.LeakyReLU(0.1), | |
| 'softplus': nn.Softplus, 'ELU': nn.ELU, 'silu': nn.SiLU} | |
| class Physics_Attention_Irregular_Mesh(nn.Module): | |
| def __init__(self, dim, heads=8, dim_head=64, dropout=0., slice_num=64): | |
| super().__init__() | |
| inner_dim = dim_head * heads | |
| self.dim_head = dim_head | |
| self.heads = heads | |
| self.scale = dim_head ** -0.5 | |
| self.softmax = nn.Softmax(dim=-1) | |
| self.dropout = nn.Dropout(dropout) | |
| self.temperature = nn.Parameter(torch.ones([1, heads, 1, 1]) * 0.5) | |
| self.in_project_x = nn.Linear(dim, inner_dim) | |
| self.in_project_fx = nn.Linear(dim, inner_dim) | |
| self.in_project_slice = nn.Linear(dim_head, slice_num) | |
| for l in [self.in_project_slice]: | |
| torch.nn.init.orthogonal_(l.weight) # use a principled initialization | |
| self.to_q = nn.Linear(dim_head, dim_head, bias=False) | |
| self.to_k = nn.Linear(dim_head, dim_head, bias=False) | |
| self.to_v = nn.Linear(dim_head, dim_head, bias=False) | |
| self.to_out = nn.Sequential( | |
| nn.Linear(inner_dim, dim), | |
| nn.Dropout(dropout) | |
| ) | |
| def forward(self, x): | |
| # B N C | |
| B, N, C = x.shape | |
| ### (1) Slice | |
| fx_mid = self.in_project_fx(x).reshape(B, N, self.heads, self.dim_head) \ | |
| .permute(0, 2, 1, 3).contiguous() # B H N C | |
| x_mid = self.in_project_x(x).reshape(B, N, self.heads, self.dim_head) \ | |
| .permute(0, 2, 1, 3).contiguous() # B H N C | |
| slice_weights = self.softmax(self.in_project_slice(x_mid) / self.temperature) # B H N G | |
| slice_norm = slice_weights.sum(2) # B H G | |
| slice_token = torch.einsum("bhnc,bhng->bhgc", fx_mid, slice_weights) | |
| slice_token = slice_token / ((slice_norm + 1e-5)[:, :, :, None].repeat(1, 1, 1, self.dim_head)) | |
| ### (2) Attention among slice tokens | |
| q_slice_token = self.to_q(slice_token) | |
| k_slice_token = self.to_k(slice_token) | |
| v_slice_token = self.to_v(slice_token) | |
| dots = torch.matmul(q_slice_token, k_slice_token.transpose(-1, -2)) * self.scale | |
| attn = self.softmax(dots) | |
| attn = self.dropout(attn) | |
| out_slice_token = torch.matmul(attn, v_slice_token) # B H G D | |
| ### (3) Deslice | |
| out_x = torch.einsum("bhgc,bhng->bhnc", out_slice_token, slice_weights) | |
| out_x = rearrange(out_x, 'b h n d -> b n (h d)') | |
| return self.to_out(out_x) | |
| class MLP(nn.Module): | |
| def __init__(self, n_input, n_hidden, n_output, n_layers=1, act='gelu', res=True): | |
| super(MLP, self).__init__() | |
| if act in ACTIVATION.keys(): | |
| act = ACTIVATION[act] | |
| else: | |
| raise NotImplementedError | |
| self.n_input = n_input | |
| self.n_hidden = n_hidden | |
| self.n_output = n_output | |
| self.n_layers = n_layers | |
| self.res = res | |
| self.linear_pre = nn.Sequential(nn.Linear(n_input, n_hidden), act()) | |
| self.linear_post = nn.Linear(n_hidden, n_output) | |
| self.linears = nn.ModuleList([nn.Sequential(nn.Linear(n_hidden, n_hidden), act()) for _ in range(n_layers)]) | |
| def forward(self, x): | |
| x = self.linear_pre(x) | |
| for i in range(self.n_layers): | |
| if self.res: | |
| x = self.linears[i](x) + x | |
| else: | |
| x = self.linears[i](x) | |
| x = self.linear_post(x) | |
| return x | |
| class Transolver_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, | |
| slice_num=32, | |
| ): | |
| super().__init__() | |
| self.last_layer = last_layer | |
| self.ln_1 = nn.LayerNorm(hidden_dim) | |
| self.Attn = Physics_Attention_Irregular_Mesh(hidden_dim, heads=num_heads, dim_head=hidden_dim // num_heads, | |
| dropout=dropout, slice_num=slice_num) | |
| 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)) + 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): | |
| def __init__(self, | |
| space_dim=1, | |
| n_layers=5, | |
| n_hidden=256, | |
| dropout=0, | |
| n_head=8, | |
| act='gelu', | |
| mlp_ratio=1, | |
| fun_dim=1, | |
| out_dim=1, | |
| slice_num=32, | |
| ref=8, | |
| unified_pos=False | |
| ): | |
| super(Model, self).__init__() | |
| self.__name__ = 'UniPDE_3D' | |
| self.ref = ref | |
| self.unified_pos = unified_pos | |
| if self.unified_pos: | |
| self.preprocess = MLP(fun_dim + self.ref * self.ref * self.ref, n_hidden * 2, n_hidden, n_layers=0, | |
| res=False, act=act) | |
| else: | |
| self.preprocess = MLP(fun_dim + space_dim, n_hidden * 2, n_hidden, n_layers=0, res=False, act=act) | |
| self.n_hidden = n_hidden | |
| self.space_dim = space_dim | |
| self.blocks = nn.ModuleList([Transolver_block(num_heads=n_head, hidden_dim=n_hidden, | |
| dropout=dropout, | |
| act=act, | |
| mlp_ratio=mlp_ratio, | |
| out_dim=out_dim, | |
| slice_num=slice_num, | |
| last_layer=(_ == n_layers - 1)) | |
| for _ in range(n_layers)]) | |
| self.initialize_weights() | |
| self.placeholder = nn.Parameter((1 / (n_hidden)) * torch.rand(n_hidden, dtype=torch.float)) | |
| 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 get_grid(self, my_pos): | |
| # my_pos 1 N 3 | |
| batchsize = my_pos.shape[0] | |
| gridx = torch.tensor(np.linspace(-1.5, 1.5, self.ref), dtype=torch.float) | |
| gridx = gridx.reshape(1, self.ref, 1, 1, 1).repeat([batchsize, 1, self.ref, self.ref, 1]) | |
| gridy = torch.tensor(np.linspace(0, 2, self.ref), dtype=torch.float) | |
| gridy = gridy.reshape(1, 1, self.ref, 1, 1).repeat([batchsize, self.ref, 1, self.ref, 1]) | |
| gridz = torch.tensor(np.linspace(-4, 4, self.ref), dtype=torch.float) | |
| gridz = gridz.reshape(1, 1, 1, self.ref, 1).repeat([batchsize, self.ref, self.ref, 1, 1]) | |
| grid_ref = torch.cat((gridx, gridy, gridz), dim=-1).cuda().reshape(batchsize, self.ref ** 3, 3) # B 4 4 4 3 | |
| pos = torch.sqrt( | |
| torch.sum((my_pos[:, :, None, :] - grid_ref[:, None, :, :]) ** 2, | |
| dim=-1)). \ | |
| reshape(batchsize, my_pos.shape[1], self.ref * self.ref * self.ref).contiguous() | |
| return pos | |
| def forward(self, data): | |
| cfd_data, geom_data = data | |
| x, fx, T = cfd_data.x, None, None | |
| x = x[None, :, :] | |
| if self.unified_pos: | |
| new_pos = self.get_grid(cfd_data.pos[None, :, :]) | |
| x = torch.cat((x, new_pos), dim=-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, :] | |
| for block in self.blocks: | |
| fx = block(fx) | |
| return fx[0] | |