File size: 4,783 Bytes
21e6c23 | 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 | import torch
import torch.nn as nn
import numpy as np
try:
from timm.layers import trunc_normal_
except ImportError:
from torch.nn.init import trunc_normal_
from onescience.modules.mlp.MLP import StandardMLP
from onescience.modules.transformer.Transolver_block import Transolver_block
class Transolver2D_plus(nn.Module):
"""
Transolver2D_plus 模型。
这是 Transolver2D 的增强版本。
主要的区别在于它使用了 'unstructured_plus' 类型的注意力机制，这种机制通常包含更复杂的切片或特征交互逻辑（如 Gumbel Softmax 采样等），以增强模型对不规则网格的捕捉能力。
Args:
space_dim (int): 空间维度。默认值: 1。
n_layers (int): 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。
"""
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(Transolver2D_plus, self).__init__()
self.__name__ = 'Transolver'
self.ref = ref
self.unified_pos = unified_pos
if self.unified_pos:
input_dim = fun_dim + space_dim + 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_plus'
)
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):
batchsize = my_pos.shape[0]
gridx = torch.tensor(np.linspace(-2, 4, self.ref), dtype=torch.float)
gridx = gridx.reshape(1, self.ref, 1, 1).repeat([batchsize, 1, self.ref, 1])
gridy = torch.tensor(np.linspace(-1.5, 1.5, self.ref), dtype=torch.float)
gridy = gridy.reshape(1, 1, self.ref, 1).repeat([batchsize, self.ref, 1, 1])
grid_ref = torch.cat((gridx, gridy), dim=-1).to(my_pos.device).reshape(batchsize, self.ref ** 2, 2)
pos = torch.sqrt(
torch.sum((my_pos[:, :, None, :] - grid_ref[:, None, :, :]) ** 2,
dim=-1)). \
reshape(batchsize, my_pos.shape[1], self.ref * self.ref).contiguous()
return pos
def forward(self, data):
x, fx, T = data.x, None, None
x = x[None, :, :]
if self.unified_pos:
new_pos = self.get_grid(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]
|