E(n) Equivariant Graph Neural Networks 项目教程

1. 项目介绍

E(n) Equivariant Graph Neural Networks（EGNNs）是一个用于学习图神经网络的新模型，该模型具有旋转、平移、反射和排列的等变性。与现有方法相比，EGNNs在中间层不需要计算昂贵的高阶表示，同时仍然能够实现竞争性或更好的性能。此外，现有方法通常局限于3维空间中的等变性，而EGNNs可以轻松扩展到更高维空间。

该项目的主要贡献包括：

• 提出了E(n) Equivariant Graph Neural Networks模型。
• 在动力系统建模、图自编码器中的表示学习和预测分子属性等方面展示了其有效性。

2. 项目快速启动

环境准备

确保你已经安装了PyTorch 1.7.1。你可以通过以下命令安装PyTorch：

pip install torch==1.7.1

克隆项目

首先，克隆项目到本地：

git clone https://github.com/vgsatorras/egnn.git
cd egnn

运行示例代码

以下是一个简单的EGNN实现示例：

import egnn_clean as eg
import torch

# 定义参数
batch_size = 8
n_nodes = 4
n_feat = 1
x_dim = 3

# 创建虚拟变量 h, x 和全连接边
h = torch.ones(batch_size * n_nodes, n_feat)
x = torch.ones(batch_size * n_nodes, x_dim)
edges, edge_attr = eg.get_edges_batch(n_nodes, batch_size)

# 初始化EGNN
egnn = eg.EGNN(in_node_nf=n_feat, hidden_nf=32, out_node_nf=1, in_edge_nf=1)

# 运行EGNN
h, x = egnn(h, x, edges, edge_attr)

3. 应用案例和最佳实践

N-body系统实验

创建N-body数据集：

cd n_body_system/dataset
python -u generate_dataset.py --num-train 10000 --seed 43 --sufix small

运行实验：

# EGNN模型
python -u main_nbody.py --exp_name exp_1_egnn_vel --model egnn_vel --max_training_samples 3000 --lr 5e-4

# GNN模型
python -u main_nbody.py --exp_name exp_1_gnn --model gnn --max_training_samples 3000 --lr 1e-3

图自编码器实验

运行图自编码器实验：

# GNN Erdos & Renyi
python -u main_ae.py --exp_name exp1_gnn_erdosrenyi --model ae --dataset erdosrenyinodes_0.25_none --K 8 --emb_nf 8 --noise_dim 0

4. 典型生态项目

相关项目

• lucidrains/egnn-pytorch: 另一个E(n)-Equivariant Graph Neural Networks的PyTorch实现，可能用于Alphafold2的复制。
• 其他图神经网络项目: 如DGL（Deep Graph Library）和PyG（PyTorch Geometric），这些项目提供了丰富的图神经网络工具和模型。

通过这些项目，你可以进一步扩展和应用E(n) Equivariant Graph Neural Networks的能力。