作者：昇腾实战派
​知识地图​：https://blog.csdn.net/Lumos_Lovegood/article/details/161601003

背景概述

Lagrangian_MGN 是 NVIDIA PhysicsNeMo 中基于图神经网络的拉格朗日粒子流体仿真模型，适用于无网格粒子型流体、固体及可变形材料的模拟。该模型将粒子视为图节点，相互作用视为边，通过自回归方式预测粒子加速度与轨迹，相比传统 CFD 方法可提速 100–1000 倍。

本文以开发者视角，详细记录在昇腾 NPU 环境下完成模型环境搭建、数据准备、训练迁移及性能验证的全过程，为同类仿真模型的 NPU 适配提供参考。

模型介绍

lagrangian_mgn 是 NVIDIA PhysicsNeMo 示例中的拉格朗日粒子流体仿真模型，基于 MeshGraphNet（MGN），用图神经网络（GNN）模拟无网格 / 粒子型流体、固体、可变形材料（如飞溅水、沙子、凝胶）。核心思想是将粒子视为图节点，相互作用视为边，通过自回归方式预测粒子加速度与轨迹。特点：无网格、自适应边界、适合大变形 / 自由面；相比传统 CFD 快 100–1000 倍。

图构建（Graph Construction）

• 节点（Node）：每个流体 / 固体粒子 → 特征向量：
  • 位置：d 维（2D/3D）
  • 历史速度：num_history × d 维（时序）
  • 节点类型：one-hot（流体 / 沙子 / 边界等）
  • 壁面特征：2×d 维（距边界距离）
• 边（Edge）：按半径邻域连接粒子 → 特征：
  • 相对位移：d 维
  • 距离：1 维

Encoder（编码器）

• 2 层 MLP，隐维 128
• 作用：节点 / 边特征 → 统一高维隐空间

Processor（处理器，核心）

• 10 层消息传递（Message-Passing）
• 每层：
  1. 边更新：MLP(节点i + 节点j + 边特征) → 新边特征
  2. 节点聚合：求和邻边特征
  3. 节点更新：MLP(原节点特征 + 聚合消息) → 新节点特征
• 隐维 128，ReLU 激活

Decoder（解码器）

• 2 层 MLP，隐维 128
• 输出：粒子加速度（d 维） → 积分得下一时刻速度 / 位置

自回归推理

• 用 t 时刻状态 预测 t+1 加速度 → 更新速度 / 位置 → 重建图 → 循环

版本信息

硬件：Atlas 800T A3训练服务器

软件	版本
cann	8.3.RC1
hdk	25.5.0
python	3.11
torch	2.7.1

CANN 镜像获取

docker pull swr.cn-south-1.myhuaweicloud.com/ascendhub/cann:8.3.rc1-a3-openeuler24.03-py3.11

• 更多版本下载请参考 CANN 镜像仓库
• 注意：若使用Atlas 800T A3服务器设备请下载带 a3 关键字的版本，Atlas 800T A2设备请下载 910b 关键字版本
• 可使用 npu-smi info 命令检查驱动是否正常

创建容器

docker run --privileged -it -u root --ipc=host --network=host \
    --device=/dev/davinci_manager \
    --device=/dev/devmm_svm \
    --device=/dev/hisi_hdc \
    -v /etc/localtime:/etc/localtime \
    -v /usr/local/Ascend/driver:/usr/local/Ascend/driver \
    -v /usr/local/sbin/npu-smi:/usr/local/sbin/npu-smi \
    -v /home:/home \
    -v /data:/data \
    --name=mgn swr.cn-south-1.myhuaweicloud.com/ascendhub/cann:8.3.rc1-a3-openeuler24.03-py3.11 /bin/bash

Conda 安装（可选）

也可使用容器系统默认 Python 环境，可忽略此步骤。

wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-aarch64.sh
bash Miniconda3-latest-Linux-aarch64.sh  # 最后初始化环境变量记得输入 yes
conda create -n mgn python=3.11
conda activate mgn

其他依赖安装

pip install -r requirements.txt

requirements.txt 内容如下：

hydra-core
omegaconf
tfrecord
wandb
warp-lang
s3fs
treelib
termcolor
cftime
einops
gitpython
h5py
jaxtyping
nvtx
onnx
pandas
tensordict
tensorboard
torch_geometric

pip install torch-scatter --no-build-isolation

torch_npu 安装

wget https://gitcode.com/Ascend/pytorch/releases/download/v7.3.1-pytorch2.7.1/torch_npu-2.7.1.post3-cp311-cp311-manylinux_2_28_aarch64.whl
pip install torch_npu-2.7.1.post3-cp311-cp311-manylinux_2_28_aarch64.whl  # torch 版本会自动安装
pip install pyyaml decorator scipy attrs psutil  # torch_npu 依赖

physicsnemo 安装

git clone https://github.com/NVIDIA/physicsnemo.git
cd physicsnemo
pip install . --no-deps

数据下载

bash examples/cfd/lagrangian_mgn/raw_datasets/download_dataset.sh Water /data/   # /data/ 为数据下载目录，可自定义

迁移配置

在启动脚本中导入 torch_npu 进行自动迁移适配。

编辑 examples/cfd/lagrangian_mgn/train.py，在文件开头添加：

import torch_npu
from torch_npu.contrib import transfer_to_npu

为了快速验证训练是否正常，我们添加了打印每个 step 训练信息的代码（仅测试阶段）：

启动训练

单卡训练命令如下：

精度测试

固定随机种子以确保可复现性：

实验条件：模型配置默认，运行 1000 个 step。实验结果：loss 收敛趋势一致，符合预期。