|
| 1 | +# e3nn 适配 paddle 设计文档 |
| 2 | + |
| 3 | +> RFC 文档相关记录信息 |
| 4 | +
|
| 5 | +| | | |
| 6 | +| ------------ | ------------------ | |
| 7 | +| 提交作者 | BeingGod | |
| 8 | +| 提交时间 | 2025-10-09 | |
| 9 | +| RFC 版本号 | v1.0 | |
| 10 | +| 依赖飞桨版本 | 3.2 版本 | |
| 11 | +| 文件名 | 20251009_e3nn | |
| 12 | + |
| 13 | +## 1. 概述 |
| 14 | + |
| 15 | +### 1.1 相关背景 |
| 16 | + |
| 17 | +e3nn(Euclidean Neural Networks) 是一个专门为 三维几何与物理问题 设计的 群等变(equivariant)深度学习库,基于 PyTorch 实现。主要用于 三维结构、物理建模、分子模拟 等需要几何等变性的场景。 |
| 18 | + |
| 19 | +### 1.2 功能目标 |
| 20 | + |
| 21 | +1. 完整实现e3nn中的全部API,精度实现对齐; |
| 22 | +2. 实现对应的单元测试; |
| 23 | + |
| 24 | + |
| 25 | +### 1.3 意义 |
| 26 | + |
| 27 | +e3nn 作为 PaddleCFD 的第三方依赖,被广泛应用于材料、流体等科学计算领域中,paddle 完整支持 e3nn 可以进一步丰富飞桨在 AI4S 领域的应用生态。 |
| 28 | + |
| 29 | +## 2. 目标调研 |
| 30 | + |
| 31 | +### 2.1 适配版本 |
| 32 | + |
| 33 | +| 组件 | 版本 | |
| 34 | +| ------------ | ------------------ | |
| 35 | +| paddle | 3.2 | |
| 36 | +| e3nn | 0.5.6 | |
| 37 | + |
| 38 | + |
| 39 | +### 2.2 相关接口及对应单测梳理 |
| 40 | + |
| 41 | +项目结构 |
| 42 | +|模块 |功能| 示例| |
| 43 | +| ------------ | ------------------ |------------------ | |
| 44 | +|e3nn.o3 |SO(3) 群操作与张量代数 |Irreps, SphericalHarmonics, TensorProduct| |
| 45 | +|e3nn.nn |等变网络层 |FullyConnectedTensorProduct, Gate, BatchNorm| |
| 46 | +|e3nn.io |输入输出封装 |转换为 IrrepsData 类型| |
| 47 | +|e3nn.util |工具函数 |坐标变换、可视化等| |
| 48 | + |
| 49 | +单测 |
| 50 | +| 单测路径 | |
| 51 | +| ------------ | |
| 52 | +| tests/math/perm_test.py| |
| 53 | +| tests/math/soft_one_hot_test.py| |
| 54 | +| tests/math/normalize_activation_test.py| |
| 55 | +| tests/math/soft_unit_step_test.py| |
| 56 | +| tests/nn/s2act_test.py| |
| 57 | +| tests/nn/activation_test.py| |
| 58 | +| tests/nn/extract_test.py| |
| 59 | +| tests/nn/gate_test.py| |
| 60 | +| tests/nn/so3act_test.py| |
| 61 | +| tests/nn/batchnorm_test.py| |
| 62 | +| tests/nn/normact_test.py| |
| 63 | +| tests/nn/fc_test.py| |
| 64 | +| tests/nn/dropout_test.py| |
| 65 | +| tests/o3/s2_test.py| |
| 66 | +| tests/o3/angular_spherical_harmonics_test.py| |
| 67 | +| tests/o3/norm_test.py| |
| 68 | +| tests/o3/irreps_test.py| |
| 69 | +| tests/o3/wigner_test.py| |
| 70 | +| tests/o3/rotation_test.py| |
| 71 | +| tests/o3/reduce_tensor_test.py| |
| 72 | +| tests/o3/tensor_product_sub_test.py| |
| 73 | +| tests/o3/cartesian_spherical_harmonics_test.py| |
| 74 | +| tests/o3/tensor_product_test.py| |
| 75 | +| tests/o3/linear_test.py| |
| 76 | +| tests/util/test_test.py| |
| 77 | +| tests/util/test_jit.py| |
| 78 | +| tests/defaults_test.py| |
| 79 | + |
| 80 | +## 3. 设计思路与实现方案 |
| 81 | + |
| 82 | +1. 使用 PaConvert 完成 python 接口迁移 |
| 83 | +2. 对不支持的接口进行手动改写 |
| 84 | +3. 跑通全量单测并实现精度对齐 |
| 85 | + |
| 86 | +### 3.1 补充说明[可选] |
| 87 | + |
| 88 | +由于 Paddle 主框架对 compile 支持有限,对于 e3nn 中涉及 torch.compile 相关的测例暂时跳过。 |
| 89 | + |
| 90 | +## 4. 测试和验收的考量 |
| 91 | + |
| 92 | +1. 实现e3nn中的API,并给出安装、使用说明; |
| 93 | +2. 实现其中对应的单元测试并通过; |
| 94 | +3. 最终代码合入PFCCLab组织下; |
| 95 | + |
| 96 | +## 5. 可行性分析和排期规划 |
| 97 | + |
| 98 | +- 2025.10:调研,复现代码并作调整。整理项目产出,撰写案例文档 |
| 99 | + |
| 100 | +## 6. 影响面 |
| 101 | + |
| 102 | +1. 在 PFCC 下添加 e3nn 仓库 |
0 commit comments