File size: 5,346 Bytes
e9b4f6f | 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 | import torch
import torch.nn as nn
import numpy as np
from timm.layers import trunc_normal_
from einops import rearrange, repeat
from onescience.modules.mlp.MLP import StandardMLP
from onescience.modules.transformer.Transolver_block import Transolver_block
class Transolver3D(nn.Module):
"""
Transolver3D 模型。
该模型专为处理三维物理场问题（如 CFD）而设计。它利用 TransolverBlock 堆叠而成的深层网络来捕捉复杂的物理依赖关系。
模型首先通过一个预处理 MLP 将输入的物理状态（和可选的统一位置编码）映射到隐空间，然后经过多层 Transolver Block 进行特征提取和交互，
最后输出预测的物理场。
Args:
space_dim (int): 空间维度（例如 1, 2, 3）。默认值: 1。
n_layers (int): Transolver Block 的层数。默认值: 5。
n_hidden (int): 隐藏层特征维度。默认值: 256。
dropout (float): Dropout 概率。默认值: 0。
n_head (int): 注意力头数。默认值: 8。
act (str): 激活函数类型。默认值: 'gelu'。
mlp_ratio (float): MLP 膨胀比率。默认值: 1。
fun_dim (int): 输入物理场的特征维度。默认值: 1。
out_dim (int): 输出特征维度。默认值: 1。
slice_num (int): 物理注意力中的切片数量。默认值: 32。
ref (int): 统一位置编码的参考网格分辨率。默认值: 8。
unified_pos (bool): 是否使用统一位置编码。默认值: False。
形状:
输入 data: 包含 x (物理状态) 和 pos (坐标) 的数据对象。
输出: (N, out_dim)，预测的物理场。
"""
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(Transolver3D, self).__init__()
self.__name__ = 'Transolver3D'
self.ref = ref
self.unified_pos = unified_pos
if self.unified_pos:
input_dim = fun_dim + self.ref * self.ref * self.ref
else:
input_dim = fun_dim + space_dim
self.preprocess = StandardMLP(
input_dim=input_dim,
hidden_dims=[n_hidden * 2],
output_dim=n_hidden,
activation=act,
use_bias=True,
use_skip_connection=False
)
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),
geotype='unstructured'
)
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).to(my_pos.device).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 = data
x, fx, T = cfd_data.x, None, None
x = x[None, :, :] # [1, N, C]
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]
|