File size: 5,181 Bytes
bdce880 | 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 | """Wrapper around the OFFICIAL Transolver++ model (thuml/Transolver_plus, ICML 2025)
so it trains/evaluates in our harness on our PyG graphs, apples to apples with GeoReNet
and the Transolver baseline (same data, split, loss, eval). We load the authors'
models/Transolver_plus.py verbatim and only adapt the I/O plus two environment shims:
1. timm: the official file does `from timm.layers import trunc_normal_`; we stub that
one symbol with torch's native trunc_normal_ (avoids the timm dependency that
conflicts with our pinned torch 2.11+cu128).
2. torch.distributed.nn.all_reduce: the eidetic attention sums slice statistics across
GPUs via all_reduce (the paper's input-invariant parallelism). On a single GPU,
all_reduce over one process is the identity, so we replace it with a passthrough.
This is the faithful single-GPU behavior of the authors' code, not a model change.
The Model.forward consumes (x, pos, condition): x = [coords + node features], pos = coords
(used only if unified_pos), condition = optional global token. We feed coords + node
features + broadcast (Re, yaw) globals so Transolver++ gets the SAME conditioning signal
GeoReNet's FiLM uses (fair), and condition=None.
"""
from __future__ import annotations
import os, sys, types, importlib.util
import torch
import torch.nn as nn
_OFFICIAL = os.path.join(os.path.dirname(os.path.abspath(__file__)),
"Transolver_plus-main", "models", "Transolver_plus.py")
def _stub_timm():
if "timm.layers" in sys.modules:
return
timm = types.ModuleType("timm")
layers = types.ModuleType("timm.layers"); layers.trunc_normal_ = torch.nn.init.trunc_normal_
models = types.ModuleType("timm.models")
mlayers = types.ModuleType("timm.models.layers"); mlayers.trunc_normal_ = torch.nn.init.trunc_normal_
timm.layers = layers; timm.models = models; models.layers = mlayers
sys.modules.update({"timm": timm, "timm.layers": layers,
"timm.models": models, "timm.models.layers": mlayers})
def _load_official():
_stub_timm()
spec = importlib.util.spec_from_file_location("transolverpp_official", _OFFICIAL)
mod = importlib.util.module_from_spec(spec)
spec.loader.exec_module(mod)
# single-GPU shims: ensure ReduceOp exists, then make all_reduce the identity
# (all_reduce over one process is a passthrough = faithful single-GPU behavior).
import torch.distributed as _td
if not hasattr(mod.dist_nn, "ReduceOp"):
mod.dist_nn.ReduceOp = getattr(_td, "ReduceOp", types.SimpleNamespace(SUM=0))
mod.dist_nn.all_reduce = lambda t, op=None: t
return mod
class _NS:
pass
class TransolverPlusWrapper(nn.Module):
"""Drop-in for our training loop: forward(data) -> [N_nodes, out_dim].
Config sized to ~GeoReNet/Transolver scale for a fair comparison; align to the
authors' main_airplane.py settings if exact faithfulness to their aero run is wanted."""
def __init__(self, node_in, space_dim=3, global_dim=2, out_dim=7, n_hidden=256,
n_layers=8, n_head=8, slice_num=64, mlp_ratio=2):
super().__init__()
mod = _load_official()
self.space_dim = space_dim
self.global_dim = global_dim
# Model preprocess input = fun_dim + space_dim; we put coords in space_dim and
# node features + broadcast globals in fun_dim.
self.net = mod.Model(space_dim=space_dim, fun_dim=node_in + global_dim, out_dim=out_dim,
n_hidden=n_hidden, n_layers=n_layers, n_head=n_head,
slice_num=slice_num, mlp_ratio=mlp_ratio, unified_pos=False)
def forward(self, data):
# PER GRAPH: a PyG DataLoader concatenates graphs into one node tensor with a
# `batch` vector. Transolver++ attention treats its input as ONE point cloud, so
# we must run each graph separately (else nodes leak across graphs and only graph
# 0's globals are used). Output is reassembled in the original node order.
N = data.x.shape[0]
batch = getattr(data, "batch", None)
if batch is None:
batch = torch.zeros(N, dtype=torch.long, device=data.x.device)
B = int(batch.max().item()) + 1
pos_all = data.pos
if pos_all.shape[1] < self.space_dim:
pos_all = torch.cat([pos_all, pos_all.new_zeros(pos_all.shape[0],
self.space_dim - pos_all.shape[1])], 1)
gf_all = data.global_feat.reshape(B, -1)[:, :self.global_dim] if self.global_dim else None
outs = []
for b in range(B):
m = batch == b
pos = pos_all[m]
parts = [pos[:, :self.space_dim], data.x[m]]
if self.global_dim:
parts.append(gf_all[b:b + 1].expand(pos.shape[0], -1))
x = torch.cat(parts, dim=-1).unsqueeze(0) # [1, n_b, space+fun]
outs.append(self.net((x, pos.unsqueeze(0), None))[0]) # [n_b, out_dim]
return torch.cat(outs, dim=0) # [N, out_dim], original order
def count_params(m):
return sum(p.numel() for p in m.parameters())
|