Transolver_plus（Transolver++）是面向百万级非结构化工业几何的 AI4S 神经 PDE 正向求解器，在初代 Transolver 基础上通过自适应物理分片注意力、分层稀疏窗口、低通信并行框架实现显存大幅下降、训练推理显著提速、复杂物理场精度提升，可用于钢材多孔结构、汽车气动、航空构件大规模仿真加速。
模型：transolver++
AI加速卡：910C
CPU架构：ARM
CANN：9.0.0
docker run -it -u root -d --net=host \
--privileged \
--ipc=host \
--device=/dev/davinci_manager \
--device=/dev/devmm_svm \
--device=/dev/hisi_hdc \
-v /usr/local/Ascend/driver:/usr/local/Ascend/driver \
-v /usr/local/dcmi:/usr/local/dcmi \
-v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
-v /usr/local/sbin:/usr/local/sbin \
-v /usr/local/Ascend/driver/tools/hccn_tool:/usr/local/Ascend/driver/tools/hccn_tool \
--name transolver \
quay.io/ascend/cann:9.0.0-a3-ubuntu22.04-py3.11 \
/bin/bash注：后续操作均是容器内进行
cd /root/
git clone https://github.com/thuml/Transolver_plus.gitapt update
apt install libxcb1 libx11-xcb1
apt install -y libgl1-mesa-glx
apt install -y libglib2.0-0 libsm6 libxext6 libxrender-dev libgomp1
apt-get install -y zlib1g-dev build-essential
pip install torch==2.10.0 torch_npu==2.10.0
pip install h5py numpy scipy pandas tqdm einops
将代码patch拷贝到目录/root/Transolver_plus/下，应用patch
cd /root/Transolver_plus
git apply transolver_plus_ascend_npu.patch主要代码修改说明如下：
1、设备适配，cuda → npu，nccl → hccl，所有 17 处 .cuda() 替换
2、DDP 修复：原始代码多卡时模型未用 DistributedDataParallel 包装，导致各 rank 独立训练、checkpoint 竞态覆盖。
已修复为正确 DDP 模式
3、推理优化：model.eval()、FP16 混合精度推理、自动混合精度上下文
4、其他工程化改进这里使用AirCraft 数据集，将数据集下载到/root/aircraft_data/（30 种飞机设计 × 5 种工况共 150 个 h5 样本（~332K 点/样本）），数据集下载好按照论文比例划分， 140/10 划分训练/测试集，与论文一致。
（注意：论文是没有说明具体划分结果的，所以不同划分下结果会有些许差异）
训练脚本内容参考 main_airplane.py
通过如下命令启动训练：
python main_airplane.py \
--data_dir /root/aircraft_data \
--save_dir /root/aircraft_output \
--nb_epochs 200 \
--lr 0.001 \
--batch_size 1 \
--device 4
训练结果如下：
官方代码中没有完整的多卡训练示例，这里通过进一步针对DDP的适配修改，可以执行多卡训练，命令参考如下：
torchrun --nproc_per_node=4 main_airplane.py \
--data_dir /root/aircraft_data \
--save_dir /root/aircraft_output_lr004 \
--lr 0.004 --batch_size 8 --nb_epochs 200
说明：这里可以采用线性缩放学习率 lr=0.004或者适当调整策略，补偿 DDP 梯度平均
推理脚本内容参考：
import sys, torch, os, warnings, numpy as np, h5py
warnings.filterwarnings("ignore")
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from dataset.dataset import AirplaneDataLoader, AirplaneDataset
from models.Transolver_plus import Model
def main():
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--data_dir', default='/root/aircraft_data')
parser.add_argument('--save_dir', default='/root/trans/aircraft_output_lr004_easy')
parser.add_argument('--ckpt', default='model_200.pth')
parser.add_argument('--device', default=0, type=int)
parser.add_argument('--save_npy', action='store_true', help='save predictions as .npy files')
args = parser.parse_args()
device = torch.device("npu", args.device)
torch.npu.set_device(device)
data_dir = args.data_dir
train_ds = AirplaneDataset(data_dir, train=True)
val_ds = AirplaneDataset(data_dir, train=False)
val_loader = AirplaneDataLoader(val_ds, batch_size=1)
all_pos, all_vals = [], []
for f in train_ds.f_list:
with h5py.File(f, "r") as h:
all_pos.append(h["pos"][:])
all_vals.append(h["values"][:])
pos = torch.from_numpy(np.concatenate(all_pos, axis=0)).float()
val = torch.from_numpy(np.concatenate(all_vals, axis=0)).float()
pm, ps = pos.mean(dim=0).view(1,1,3).npu(), pos.std(dim=0).view(1,1,3).npu()
om, os_ = val.mean(dim=0).view(1,1,-1).npu(), val.std(dim=0).view(1,1,-1).npu()
model = Model(n_hidden=256, n_layers=4, space_dim=7, fun_dim=0, n_head=8,
mlp_ratio=2, out_dim=6, slice_num=32, unified_pos=0, dropout=0.1).npu()
sd = torch.load(os.path.join(args.save_dir, args.ckpt), map_location="npu", weights_only=True)
# handle DDP wrapped checkpoint
if any(k.startswith("module.") for k in sd):
sd = {k.removeprefix("module."): v for k, v in sd.items()}
model.load_state_dict(sd)
model.eval()
ch_names = ["Cp", "rho", "u(CL)", "v", "w", "ch5(SURF)"]
l2res = {c: [] for c in range(6)}
fnames = []
with torch.no_grad():
for i, (x, y, pos, geom, _) in enumerate(val_loader):
fname = os.path.basename(val_ds.f_list[i])
fnames.append(fname)
x, pos, y, geom = x.npu(), pos.npu(), y.npu(), geom.npu()
pos_n = (pos - pm) / ps
x[:,:,:3] = pos_n
out = model((x, pos_n, geom)) * os_ + om
for c in range(6):
l2r = torch.norm(out[0,:,c] - y[0,:,c]).item() / torch.norm(y[0,:,c]).item()
l2res[c].append(l2r)
if args.save_npy:
stem = os.path.splitext(fname)[0]
np.save(os.path.join(args.save_dir, f"{stem}.npy"), out.cpu().numpy())
print(f"\nPer-sample L2RE (ch5):")
for i, fn in enumerate(fnames):
print(f" {fn}: {l2res[5][i]:.4f}")
print(f"\nPer-channel L2RE:")
for c in range(6):
print(f" ch{c} ({ch_names[c]:>10s}): {np.mean(l2res[c]):.4f}")
print(f"\n SURF(ch5) = {np.mean(l2res[5]):.4f}")
print(f" CL(ch2) = {np.mean(l2res[2]):.4f}")
if __name__ == '__main__':
main()
也可以参考官方的推理脚本（依赖 result.csv）
命令参考：
python eval.py \
--data_dir /root/aircraft_data \
--save_dir /path/to/checkpoint_dir \
--device 0执行结果：
1、精度对齐
使用论文中同样的数据集，最终验证的精度结果:
SURF = 0.064
CL = 0.012
与论文中的精度效果一致。
2、单卡容量验证
Transolver++ 通过消除冗余 f 投影），单卡可处理 1.2M 点（A100 40GB）。在 Ascend NPU 910C 上通过合成数据验证此声称。
验证方案：
使用 capacity_test.py 生成椭球面合成网格数据（pos + normals + values），匹配 AirCraft 的 h5 格式
在不同点数规模下加载模型（4 层, hidden=256, 8 head, slice_num=32，与 AirCraft 配置一致），分别测量推理（forward only）和训练（forward + backward + optim.step）的峰值显存。峰值显存通过 torch.npu.reset_peak_memory_stats\) + torch.npu.max\_memory\_allocated\( 采集。
验证结果： NPU 910C单卡（ 65.8GB）训练模式下可处理约 1.9M 点，按照显存比计算与官方宣称效果基本一致。