File size: 3,716 Bytes
3e77c56 | 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 | """Physics-Attention for irregular meshes / point clouds.
This module is reused **verbatim** from the Transolver project and is licensed under the MIT
License. Attribution preserved per the license terms.
Transolver: A Fast Transformer Solver for PDEs on General Geometries
Wu, Luo, Wang, Wang, Long (THUML / Tsinghua University), ICML 2024 Spotlight.
Source: https://github.com/thuml/Transolver (MIT License)
Paper: arXiv:2402.02366
Key constants preserved exactly: temperature init 0.5, orthogonal init of the slice
projection, slice-norm epsilon 1e-5, attention scale dim_head**-0.5.
Note: in the assembled model (see transolver.py), the operator's global ``trunc_normal_``
weight initialization runs *after* this block is constructed and overwrites the orthogonal
init of ``in_project_slice`` — this matches the upstream assembly order and is what produced
the published 0.0064. We keep the orthogonal init here to stay verbatim with upstream.
"""
import torch
import torch.nn as nn
from einops import rearrange
class Physics_Attention_Irregular_Mesh(nn.Module):
## for irregular meshes in 1D, 2D or 3D space
# NOTE: the dim_head=64 default is upstream's generic default. For the Elasticity config the
# assembler (transolver.py) always passes dim_head = n_hidden // n_heads = 16; the default here
# is never used in this repo. See docs/RECONCILIATION.md §3.
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)
|