【几何深度学习（Geometric Deep Learning）】物理对称性嵌入网络

一个 可直接运行 的完整示例，演示如何把物理对称性（这里是 3‑D 空间中的 SO(3) 旋转）嵌入到深度网络里。
我们使用的是 e3nn（一个专门为 SE(3)-equivariant 网络提供的 PyTorch 生态），并结合 torch_geometric 来处理分子图数据。

⚠️ 前置条件

• Python ≥ 3.9
• CUDA (可选，若无 GPU 可改用 CPU)
• 已安装 pip install torch e3nn torch-geometric rdkit tqdm

---

1️⃣ 数据准备

我们以公开的 QM9 数据集为例（分子能量预测）。torch_geometric 自带下载工具。

# qm9_data.py
import os
from torch_geometric.datasets import QM9
from torch_geometric.loader import DataLoader

def get_qm9_loader(batch_size=128, split='train'):
    root = './data/qm9'
    dataset = QM9(root)

    if split == 'train':
        data_list = dataset[:13330]