在cu12.8上安装Nvidia PhysicsNemo

Python版本要求

系统要求： Python ≥ 3.10

Pytorch安装

安装PhysicsNemo后需要使用graph，所以要安装dgl。dgl目前官方支持PyTorch最高版本为2.4.x+cu12.4，因此先安装此版本的PyTorch。

conda install pytorch==2.4.1 torchvision==0.19.1 torchaudio==2.4.1 pytorch-cuda=12.4 -c pytorch -c nvidia

PhysicsNemo安装

使用pip安装（经测试可以更换成tsinghua镜像源）

pip install nvidia-physicsnemo
pip install Cython
pip install nvidia-physicsnemo-sym --no-build-isolation

安装完成后使用以下代码进行检测

import torch
from physicsnemo.models.mlp.fully_connected import FullyConnected
model = FullyConnected(in_features=32, out_features=64)
input = torch.randn(128, 32)
output = model(input)
output.shape
>> torch.Size([128, 64])

dgl安装

https://www.dgl.ai/pages/start.html
dgl目前官方支持PyTorch最高版本为2.4.x+cu12.4

如果只安装了PhysicsNemo，运行graph相关代码时会遇到

No module named 'einops' >> pip install einops

conda install -c dglteam/label/th24_cu124 dgl >> conda install -c dglteam/label/th24_cu124 dgl

可能会遇到

An HTTP error occurred when trying to retrieve this URL.
HTTP errors are often intermittent, and a simple retry will get you on your way.

重新运行一下即可。