PyTorch 版 SE(3)-Transformers

本仓库提供了用于训练 SE(3)-Transformer 模型以实现最先进精度的脚本和方案。仓库内容由 NVIDIA 测试并维护。

模型概述

SE(3)-Transformer 是一种图神经网络，它采用了自注意力的变体来处理 3D 点和图数据。该模型在连续 3D 旋转平移变换下具有等变性，这意味着当输入（图或点集）在 3D 空间中旋转（或更一般地说，经历正常刚性变换）时，模型输出要么保持不变，要么随输入一起变换。等变性的数学保证非常重要，它能确保在数据输入存在干扰性变换时，以及当问题具有我们希望利用的某些固有对称性时，模型性能稳定且可预测。

该模型基于以下 publications：

SE(3)-Transformers: 3D Roto-Translation Equivariant Attention Networks（NeurIPS 2020），作者：Fabian B. Fuchs、Daniel E. Worrall 等人。
Tensor field networks: Rotation- and translation-equivariant neural networks for 3D point clouds，作者：Nathaniel Thomas、Tess Smidt 等人。

后续的一篇论文阐述了如何迭代使用该模型，例如用于预测或优化蛋白质结构：

Iterative SE(3)-Transformers，作者：Fabian B. Fuchs、Daniel E. Worrall 等人。

与官方实现一样，本实现使用 PyTorch 和深度图库（DGL）。

本 SE(3)-Transformer 实现与官方实现之间的主要区别如下：

支持多 GPU 训练和推理
支持混合精度训练和推理
使用 DGL 提供的 QM9 数据集，并会自动下载
显著提高了吞吐量
显著降低了内存消耗
可选择在全连接径向轮廓层中使用层归一化（--use_layer_norm），默认关闭
可选择在注意力层之间使用等变归一化（--norm），默认关闭
用于计算基矩阵的球谐函数和 Clebsch–Gordan 系数通过 e3nn 库计算

该模型能够预测 QM9 数据集中小有机分子的量子化学性质。在这种情况下，所利用的对称性是这些性质不依赖于分子在空间中的取向或位置。

该模型使用 NVIDIA Volta、NVIDIA Turing 和 NVIDIA Ampere GPU 架构上的 Tensor Cores 进行混合精度训练。因此，研究人员可以获得比不使用 Tensor Cores 的训练快高达 1.5 倍的结果，同时还能享受混合精度训练带来的好处。该模型会针对每个 NGC 月度容器版本进行测试，以确保长期的精度和性能一致性。

模型架构

该模型由堆叠的等变图自注意力层和等变归一化层组成。最后，应用张量场网络（Tensor Field Network）卷积以获取不变特征。对这些特征执行图池化（对节点进行平均或最大池化），并将结果输入最终的多层感知机（MLP）以得到标量预测。

在此设置下，该模型是一个图到标量的网络。移除池化层可得到图到图的网络，并且可以修改最终的 TFN 以输出任何类型的特征（不变标量、三维向量等）。

模型高层架构

默认配置

SE(3)-Transformers 引入了一种用于图的自注意力层，该层对三维旋转平移具有等变性。它通过利用张量场网络（Tensor Field Networks）来构建不变的注意力权重和等变的注意力值，从而实现这一点。将等变值与不变权重相结合，产生等变输出。借助对等变特征范数进行操作的等变归一化层，该输出在归一化的同时保持了等变性。

该模型实现了以下功能：

支持任意阶数的边特征（1D、3D、5D 等），而官方实现仅支持标量不变边特征（0 阶）。阶数大于 1 的边特征会被拼接至相同阶数的节点特征。这对于复现点云处理领域已发表的结果是必需的。
数据并行多 GPU 训练（DDP）
混合精度训练（自动类型转换、梯度缩放）
梯度累积
模型 checkpoint

该模型实现了以下性能优化：

通用优化

提供了在训练开始时预计算基函数的选项，而非在每个前向传播开始时计算（--precompute_bases）
使用 torch.jit.script 对基函数计算进行即时（JIT）编译
克莱布施 - 戈登系数（Clebsch-Gordon coefficients）缓存在内存中

张量场网络优化

每个径向轮廓网络的最后一层不添加任何偏置，以避免大型广播操作
优化基函数张量的布局（维度顺序），以避免在下游 TFN 层中复制到连续内存
当 Tensor Cores 可用且计算基函数的输出特征维度为奇数时，会用零填充该维度，以更有效地利用 Tensor Cores（AMP 和 TF32 精度）
提供 TFN 卷积（和径向轮廓）的多级融合，并在条件满足时自动使用
提供低内存模式，该模式会以吞吐量为代价减少内存使用（--low_memory）

自注意力优化

在每个注意力层中，注意力键（keys）和值（values）通过单次部分 TFN 图卷积计算得出，而非两次
如果条件满足，不同输出阶数的图操作可以融合在一起

归一化优化

当满足特定条件时，等变归一化层从多个层归一化优化为基于融合范数的组归一化

针对以下超参数（与原始文献中的超参数完全相同），提供了具有竞争力的训练结果和分析：

层数：7
阶数：4
通道数：32
注意力头数：8
通道分割数：2
使用等变归一化：true
使用层归一化：true
池化方式：max

功能支持矩阵

该模型支持以下功能：

功能	SE(3)-Transformer
Automatic mixed precision (AMP)	是
Distributed data parallel (DDP)	是

功能说明

分布式数据并行（DDP）

DistributedDataParallel (DDP) 在模块级别实现数据并行，可在多个 GPU 或机器上运行。

自动混合精度（AMP）

本实现采用 PyTorch 原生的 AMP 混合精度训练方案。只需修改少量代码，就能实现 FP16 训练与 FP32 主权重存储的结合。下一节将详细解释混合精度的原理。

混合精度训练

混合精度是指在计算方法中结合使用不同的数值精度。混合精度训练通过在半精度格式下执行运算，同时在单精度中存储关键信息以保留网络关键部分的尽可能多信息，从而显著提升计算速度。自从 NVIDIA Volta 架构引入 Tensor Cores，并在后续的 NVIDIA Turing 和 NVIDIA Ampere 架构中得到发展，切换到混合精度训练可带来显著的速度提升——在计算密集型最强的模型架构上，整体速度提升最高可达 3 倍。以前使用混合精度训练需要两个步骤：

适当将模型移植为使用 FP16 数据类型。
添加损失缩放以保留小梯度值。

AMP 可在 NVIDIA Volta、NVIDIA Turing 和 NVIDIA Ampere GPU 架构上自动启用混合精度训练。PyTorch 框架代码会在内部进行所有必要的模型修改。

有关以下方面的信息：

如何使用混合精度进行训练，请参考混合精度训练论文和混合精度训练文档。
混合精度训练所使用的技术，请参考深度神经网络的混合精度训练博客。
用于混合精度训练的 APEX 工具，请参考 NVIDIA Apex：PyTorch 中轻松实现混合精度训练的工具。

启用混合精度

在 PyTorch 中，通过使用原生的自动混合精度包启用混合精度。该包在获取变量时将其转换为半精度，同时以单精度格式存储变量。此外，为了在反向传播中保留小梯度幅度，应用梯度时必须包含损失缩放步骤。在 PyTorch 中，可使用 GradScaler 自动应用损失缩放。自动混合精度在 PyTorch 内部进行所有调整，与手动操作相比有两个优势。首先，程序员无需修改网络模型代码，减少了开发和维护工作。其次，使用 AMP 可保持与所有 PyTorch 模型定义和运行 API 的向前和向后兼容性。

要启用混合精度，只需在运行训练或推理脚本时使用 --amp 标志。

启用 TF32

TensorFloat-32（TF32）是 NVIDIA A100 GPU 中用于处理矩阵数学（也称为张量运算）的新数学模式。在 A100 GPU 的 Tensor Cores 上运行 TF32，与 NVIDIA Volta GPU 上的单精度浮点数学（FP32）相比，可提供高达 10 倍的速度提升。

TF32 Tensor Cores 可以加速使用 FP32 的网络，通常不会损失精度。对于权重或激活需要高动态范围的模型，它比 FP16 更稳健。

有关更多信息，请参考 A100 GPU 中的 TensorFloat-32 将 AI 训练、HPC 加速高达 20 倍博客文章。

TF32 在 NVIDIA Ampere GPU 架构中受支持，并且默认启用。

术语表

Degree (type) - 阶数（类型）

在模型中，每个特征（输入、输出和隐藏特征）都相对于输入图以等变方式进行变换。当我们定义一个特征时，除了通道数之外，还需要选择它所遵循的变换规则。

特征的阶数或类型是一个正整数，用于描述当输入在三维空间中旋转时该特征的变换方式。

这与不同旋转阶数的不可约表示相关。

特征的阶数决定了其维度。一个d阶特征的维度为2d+1。

一些常见示例包括：

0阶：1D标量，对旋转具有不变性
1阶：3D向量，根据3D旋转矩阵进行旋转
2阶：5D向量，根据5DWigner-D矩阵进行旋转。这些可以表示对称无迹3x3矩阵。

Fiber - 纤维

纤维可以看作是一组不同类型或阶数（正整数）特征的表示，其中每种特征类型都按照其自身规则进行变换。

在本仓库中，纤维可以视为一个字典，其中键为阶数，值为通道数。

Multiplicity - 重数

给定类型特征的重数是指该特征的通道数量。

Tensor Field Network - 张量场网络

张量场网络是一种等变图卷积，它能够组合不同阶数的特征并生成新的特征，同时借助张量积保持等变性。

Equivariance - 等变性

等变性是模型函数的一种性质，指的是对输入应用对称变换后再计算函数，与先计算函数再对输出应用该变换，得到的结果相同。

在SE(3)-Transformer中，对称群是连续旋转平移群（SE(3)）。

环境设置

以下部分列出了开始训练 SE(3)-Transformer 模型所需满足的条件。

环境要求

本仓库包含一个 Dockerfile，它基于 PyTorch 21.07 NGC 容器进行扩展，并封装了部分依赖项。除了这些依赖项外，请确保您拥有以下组件：

NVIDIA Docker
PyTorch 21.07+ NGC 容器
支持的 GPU：

有关如何开始使用 NGC 容器的更多信息，请参考 NVIDIA GPU Cloud 文档和深度学习文档中的以下部分：

对于无法使用 PyTorch NGC 容器来设置所需环境或创建自己容器的用户，请参考版本化的 NVIDIA 容器支持矩阵。

快速入门指南

要使用 Tensor Cores 的混合精度或 TF32 精度，或 FP32 精度训练模型，请按照以下步骤，使用 SE(3)-Transformer 模型的默认参数在 QM9 数据集上进行操作。有关训练和推理的具体细节，请参考高级部分。

克隆仓库。

git clone https://github.com/NVIDIA/DeepLearningExamples
cd DeepLearningExamples/PyTorch/DrugDiscovery/SE3Transformer

构建 se3-transformer PyTorch NGC 容器。
```
docker build -t se3-transformer .
```

在 NGC 容器中启动交互式会话以运行训练/推理。

mkdir -p results
docker run -it --runtime=nvidia --shm-size=8g --ulimit memlock=-1 --ulimit stack=67108864 --rm -v ${PWD}/results:/results se3-transformer:latest

开始训练。
```
bash scripts/train.sh
```
开始推理/预测。
```
bash scripts/predict.sh
```

现在您已经训练并评估了模型，可以选择将您的训练结果与我们的训练精度结果进行比较。您也可以选择对性能进行基准测试，以对比训练性能基准或推理性能基准。按照这些部分中的步骤操作，将确保您获得与结果部分中所述相同的精度和性能结果。

高级说明

以下部分将详细介绍数据集、训练与推理的运行方式以及训练结果。

脚本与示例代码

在根目录下，最重要的文件包括：

Dockerfile：包含运行 SE(3)-Transformers 所需基本依赖的容器
requirements.txt：运行 SE(3)-Transformers 所需的额外依赖集合
se3_transformer/data_loading/qm9.py：QM9 数据加载、预处理以及基矩阵预计算
se3_transformer/model/layers/：包含模型架构各层的目录
se3_transformer/model/transformer.py：Transformer 主模块
se3_transformer/model/basis.py：基矩阵计算逻辑
se3_transformer/runtime/training.py：训练脚本，需作为 Python 模块运行
se3_transformer/runtime/inference.py：推理脚本，需作为 Python 模块运行
se3_transformer/runtime/metrics.py：支持多 GPU 同步的 MAE 指标
se3_transformer/runtime/loggers.py：DLLogger 和 W&B 日志记录器

参数

training.py 脚本的完整可用参数列表如下：

通用参数

--epochs：训练轮数（默认：单 GPU 下为 100）
--batch_size：批处理大小（默认：240）
--seed：全局随机种子（默认：None）
--num_workers：数据加载工作进程数（默认：8）
--amp：使用自动混合精度（默认：false）
--gradient_clip：梯度范数裁剪（默认：None）
--accumulate_grad_batches：梯度累积（默认：1）
--ckpt_interval：每 N 轮保存一次检查点（默认：-1）
--eval_interval：每 N 轮进行一次评估（默认：1）
--silent：最小化标准输出（默认：false）

路径参数

--data_dir：数据存放或下载目录（默认：./data）
--log_dir：结果日志保存目录（默认：/results）
--save_ckpt_path：检查点保存文件路径（默认：None）
--load_ckpt_path：待加载的检查点文件路径（默认：None）

优化器参数

--optimizer：使用的优化器（默认：adam）
--learning_rate：学习率（默认：单 GPU 下为 0.002）
--momentum：动量（默认：0.9）
--weight_decay：权重衰减（默认：0.1）

QM9 数据集参数

--task：训练的回归任务（默认：homo）
--precompute_bases：在脚本开始时、数据集初始化阶段预计算基矩阵，而非在每次前向传播开始时计算（默认：false）

模型架构参数

--num_layers：堆叠的 Transformer 层数（默认：7）
--num_heads：自注意力头数（默认：8）
--channels_div：送入注意力层前的通道数分割比例（默认：2）
--pooling：图池化类型（默认：max）
--norm：在每个注意力块后应用归一化层（默认：false）
--use_layer_norm：在 MLP 层之间应用层归一化（默认：false）
--low_memory：若为 true，将使用融合操作，速度较慢但内存占用更少（预计减少 25% 内存）。仅在 NVIDIA Volta GPU 上启用 AMP 或在 Ampere GPU 上运行时生效（默认：false）
--num_degrees：使用的度数数量。隐藏特征将具有类型 [0, ..., num_degrees - 1]（默认：4）
--num_channels：隐藏特征的通道数（默认：32）

命令行选项

要查看所有可用选项及其说明的完整列表，请使用 -h 或 --help 命令行选项，例如：python -m se3_transformer.runtime.training --help。

数据集指南

演示数据集

SE(3)-Transformer 是在 QM9 数据集上进行训练的。

QM9 数据集托管在 DGL 服务器上，在需要时会自动下载（38MB）。默认情况下，它存储在 ./data 目录中，但可以通过 --data_dir 参数更改此位置。

该数据集保存为 qm9_edge.npz 文件，并在运行时转换为 DGL 图。

我们使用的输入特征如下：

节点特征（6 维）：
- 原子类型的独热编码（5 维）（原子类型：H、C、N、O、F）
- 每个原子的质子数（1 维）
边特征：键类型的独热编码（4 维）（键类型：单键、双键、三键、芳香键）
相邻节点（原子）之间的相对位置

自定义数据集

要在新数据集上使用此网络，您可以扩展 se3_transformer/data_loading/data_module.py 中提供的 DataModule 类。

您的自定义 collate 函数应返回一个元组，其中包含：

一个（批处理的）DGLGraph 对象
节点特征字典（{'{degree}': tensor}）
边特征字典（{'{degree}': tensor}）
（可选）预计算的基，作为一个字典
标签，作为一个张量

然后，您可以修改 training.py 和 inference.py 脚本以使用您的新数据模块。

训练过程

训练脚本为 se3_transformer/runtime/training.py，需作为模块运行：python -m se3_transformer.runtime.training。

日志

默认情况下，生成的日志存储在 /results/ 中。可以通过 --log_dir 更改此位置。

您可以通过设置 WANDB_API_KEY 环境变量来连接现有的 Weights & Biases 账户。

检查点

可以通过 --save_ckpt_path 参数设置保存检查点的文件路径。还可以通过 --ckpt_interval 设置检查点之间的间隔（以 epoch 数为单位）。

评估

评估指标为平均绝对误差（MAE）。

--eval_interval 可以设置评估轮次之间的间隔（以 epoch 数为单位）。默认情况下，每个 epoch 后执行一次评估。

自动混合精度

要启用混合精度训练，请添加 --amp 标志。

多 GPU 和多节点

训练脚本支持 PyTorch elastic 启动器，以在多个 GPU 或节点上运行。请参考官方文档。

例如，要使用 AMP 在所有可用 GPU 上进行训练：

python -m torch.distributed.run --nnodes=1 --nproc_per_node=gpu --module se3_transformer.runtime.training --amp

推理流程

可通过 se3_transformer.runtime.inference Python 模块运行推理。

推理脚本为 se3_transformer/runtime/inference.py，需以模块形式运行：python -m se3_transformer.runtime.inference。该脚本需要一个预训练模型检查点（通过 --load_ckpt_path 参数传入）。

性能

本文档中的性能测量是在发布时进行的，可能无法反映 NVIDIA 最新软件版本所实现的性能。有关最新的性能测量结果，请访问 NVIDIA 数据中心深度学习产品性能。

基准测试

以下部分介绍如何运行基准测试以衡量模型在训练和推理模式下的性能。

训练性能基准测试

要在特定批处理大小下对训练性能进行基准测试，单 GPU 运行 bash scripts/benchmarck_train.sh {BATCH_SIZE}，多 GPU 运行 bash scripts/benchmarck_train_multi_gpu.sh {BATCH_SIZE}。

推理性能基准测试

要在特定批处理大小下对推理性能进行基准测试，运行 bash scripts/benchmarck_inference.sh {BATCH_SIZE}。

结果

以下部分详细介绍我们在训练和推理中如何实现性能和精度。

训练精度结果

训练精度：NVIDIA DGX A100（8x A100 80GB）

我们的结果是通过在配备 8x A100 80GB GPU 的 NVIDIA DGX A100 上，在 PyTorch 21.07 NGC 容器中运行 scripts/train.sh 训练脚本获得的。

GPU 数量	每 GPU 批处理大小	TF32 绝对误差	混合精度绝对误差	TF32 训练时间	混合精度训练时间	混合精度相对 TF32 的训练速度提升
1	240	0.03456	0.03460	1小时23分钟	1小时03分钟	1.32倍
8	240	0.03417	0.03424	15分钟	12分钟	1.25倍

训练精度：NVIDIA DGX-1（8x V100 16GB）

我们的结果是通过在配备 8x V100 16GB GPU 的 NVIDIA DGX-1 上，在 PyTorch 21.07 NGC 容器中运行 scripts/train.sh 训练脚本获得的。

GPU 数量	每 GPU 批处理大小	FP32 绝对误差	混合精度绝对误差	FP32 训练时间	混合精度训练时间	混合精度相对 FP32 的训练速度提升
1	240	0.03432	0.03439	2小时25分钟	1小时33分钟	1.56倍
8	240	0.03380	0.03495	29分钟	20分钟	1.45倍

训练性能结果

训练性能：NVIDIA DGX A100（8x A100 80GB）

我们的结果是通过在配备 8x A100 80GB GPU 的 NVIDIA DGX A100 上，在 PyTorch 21.07 NGC 容器中运行 scripts/benchmark_train.sh 和 scripts/benchmark_train_multi_gpu.sh 基准测试脚本获得的。性能数据（以分子/毫秒为单位）是在一个预热 epoch 后，对五个完整训练 epoch 取平均值得到的。

GPU 数量	每 GPU 批处理大小	TF32 吞吐量 [mol/ms]	混合精度吞吐量 [mol/ms]	混合精度相对 TF32 的吞吐量提升	TF32 弱扩展性	混合精度弱扩展性
1	240	2.21	2.92	1.32倍
1	120	1.81	2.04	1.13倍
8	240	17.15	22.95	1.34倍	7.76	7.86
8	120	13.89	15.62	1.12倍	7.67	7.66

要获得相同结果，请按照快速入门指南中的步骤操作。

训练性能：NVIDIA DGX-1（8x V100 16GB）

我们的结果是通过在配备 8x V100 16GB GPU 的 NVIDIA DGX-1 上，在 PyTorch 21.07 NGC 容器中运行 scripts/benchmark_train.sh 和 scripts/benchmark_train_multi_gpu.sh 基准测试脚本获得的。性能数据（以分子/毫秒为单位）是在一个预热 epoch 后，对五个完整训练 epoch 取平均值得到的。

GPU 数量	每 GPU 批处理大小	FP32 吞吐量 [mol/ms]	混合精度吞吐量 [mol/ms]	混合精度相对 FP32 的吞吐量提升	FP32 弱扩展性	混合精度弱扩展性
1	240	1.25	1.88	1.50倍
1	120	1.03	1.41	1.37倍
8	240	9.33	14.02	1.50倍	7.46	7.46
8	120	7.39	9.41	1.27倍	7.17	6.67

要获得相同结果，请按照快速入门指南中的步骤操作。

推理性能结果

推理性能：NVIDIA DGX A100（1x A100 80GB）

我们的结果是通过在配备 1x A100 80GB GPU 的 NVIDIA DGX A100 上，在 PyTorch 21.07 NGC 容器中运行 scripts/benchmark_inference.sh 推理基准测试脚本获得的。

FP16

批处理大小	平均吞吐量 [mol/ms]	平均延迟 [ms]	90% 延迟 [ms]	95% 延迟 [ms]	99% 延迟 [ms]
1600	11.60	140.94	138.29	140.12	386.40
800	10.74	75.69	75.74	76.50	79.77
400	8.86	45.57	46.11	46.60	49.97

TF32

批处理大小	平均吞吐量 [mol/ms]	平均延迟 [ms]	90% 延迟 [ms]	95% 延迟 [ms]	99% 延迟 [ms]
1600	8.58	189.20	186.39	187.71	420.28
800	8.28	97.56	97.20	97.73	101.13
400	7.55	53.38	53.72	54.48	56.62

要获得相同结果，请按照快速入门指南中的步骤操作。

推理性能：NVIDIA DGX-1（1x V100 16GB）

我们的结果是通过在配备 1x V100 16GB GPU 的 NVIDIA DGX-1 上，在 PyTorch 21.07 NGC 容器中运行 scripts/benchmark_inference.sh 推理基准测试脚本获得的。

FP16

批处理大小	平均吞吐量 [mol/ms]	平均延迟 [ms]	90% 延迟 [ms]	95% 延迟 [ms]	99% 延迟 [ms]
1600	6.42	254.54	247.97	249.29	721.15
800	6.13	132.07	131.90	132.70	140.15
400	5.37	75.12	76.01	76.66	79.90

FP32

批处理大小	平均吞吐量 [mol/ms]	平均延迟 [ms]	90% 延迟 [ms]	95% 延迟 [ms]	99% 延迟 [ms]
1600	3.39	475.86	473.82	475.64	891.18
800	3.36	239.17	240.64	241.65	243.70
400	3.17	126.67	128.19	128.82	130.54

要获得相同结果，请按照快速入门指南中的步骤操作。

发布说明

更新日志

2021年8月

初始版本发布

已知问题

如果在创建 Dataloader 迭代器期间（更准确地说，在 fork() 过程中）遇到 OSError: [Errno 12] Cannot allocate memory 错误，这很可能是由于使用了 --precompute_bases 标志。如果无法为您的机器添加更多 RAM 或 Swap 空间，建议通过移除 --precompute_bases 标志或使用 --precompute_bases false 来关闭基向量预计算功能。

PyTorch 版 SE(3)-Transformers

本仓库提供了用于训练 SE(3)-Transformer 模型以实现最先进精度的脚本和方案。仓库内容由 NVIDIA 测试并维护。

模型概述

该模型基于以下 publications：

SE(3)-Transformers: 3D Roto-Translation Equivariant Attention Networks（NeurIPS 2020），作者：Fabian B. Fuchs、Daniel E. Worrall 等人。
Tensor field networks: Rotation- and translation-equivariant neural networks for 3D point clouds，作者：Nathaniel Thomas、Tess Smidt 等人。

后续的一篇论文阐述了如何迭代使用该模型，例如用于预测或优化蛋白质结构：

Iterative SE(3)-Transformers，作者：Fabian B. Fuchs、Daniel E. Worrall 等人。

与官方实现一样，本实现使用 PyTorch 和深度图库（DGL）。

本 SE(3)-Transformer 实现与官方实现之间的主要区别如下：

支持多 GPU 训练和推理
支持混合精度训练和推理
使用 DGL 提供的 QM9 数据集，并会自动下载
显著提高了吞吐量
显著降低了内存消耗
可选择在全连接径向轮廓层中使用层归一化（--use_layer_norm），默认关闭
可选择在注意力层之间使用等变归一化（--norm），默认关闭
用于计算基矩阵的球谐函数和 Clebsch–Gordan 系数通过 e3nn 库计算

该模型能够预测 QM9 数据集中小有机分子的量子化学性质。在这种情况下，所利用的对称性是这些性质不依赖于分子在空间中的取向或位置。

模型架构

模型高层架构

默认配置

该模型实现了以下功能：

支持任意阶数的边特征（1D、3D、5D 等），而官方实现仅支持标量不变边特征（0 阶）。阶数大于 1 的边特征会被拼接至相同阶数的节点特征。这对于复现点云处理领域已发表的结果是必需的。
数据并行多 GPU 训练（DDP）
混合精度训练（自动类型转换、梯度缩放）
梯度累积
模型 checkpoint

该模型实现了以下性能优化：

通用优化

提供了在训练开始时预计算基函数的选项，而非在每个前向传播开始时计算（--precompute_bases）
使用 torch.jit.script 对基函数计算进行即时（JIT）编译
克莱布施 - 戈登系数（Clebsch-Gordon coefficients）缓存在内存中

张量场网络优化

每个径向轮廓网络的最后一层不添加任何偏置，以避免大型广播操作
优化基函数张量的布局（维度顺序），以避免在下游 TFN 层中复制到连续内存
当 Tensor Cores 可用且计算基函数的输出特征维度为奇数时，会用零填充该维度，以更有效地利用 Tensor Cores（AMP 和 TF32 精度）
提供 TFN 卷积（和径向轮廓）的多级融合，并在条件满足时自动使用
提供低内存模式，该模式会以吞吐量为代价减少内存使用（--low_memory）

自注意力优化

在每个注意力层中，注意力键（keys）和值（values）通过单次部分 TFN 图卷积计算得出，而非两次
如果条件满足，不同输出阶数的图操作可以融合在一起

归一化优化

当满足特定条件时，等变归一化层从多个层归一化优化为基于融合范数的组归一化

针对以下超参数（与原始文献中的超参数完全相同），提供了具有竞争力的训练结果和分析：

层数：7
阶数：4
通道数：32
注意力头数：8
通道分割数：2
使用等变归一化：true
使用层归一化：true
池化方式：max

功能支持矩阵

该模型支持以下功能：

功能	SE(3)-Transformer
Automatic mixed precision (AMP)	是
Distributed data parallel (DDP)	是

功能说明

分布式数据并行（DDP）

DistributedDataParallel (DDP) 在模块级别实现数据并行，可在多个 GPU 或机器上运行。

自动混合精度（AMP）

本实现采用 PyTorch 原生的 AMP 混合精度训练方案。只需修改少量代码，就能实现 FP16 训练与 FP32 主权重存储的结合。下一节将详细解释混合精度的原理。

混合精度训练

适当将模型移植为使用 FP16 数据类型。
添加损失缩放以保留小梯度值。

AMP 可在 NVIDIA Volta、NVIDIA Turing 和 NVIDIA Ampere GPU 架构上自动启用混合精度训练。PyTorch 框架代码会在内部进行所有必要的模型修改。

有关以下方面的信息：

如何使用混合精度进行训练，请参考混合精度训练论文和混合精度训练文档。
混合精度训练所使用的技术，请参考深度神经网络的混合精度训练博客。
用于混合精度训练的 APEX 工具，请参考 NVIDIA Apex：PyTorch 中轻松实现混合精度训练的工具。

启用混合精度

要启用混合精度，只需在运行训练或推理脚本时使用 --amp 标志。

启用 TF32

TF32 Tensor Cores 可以加速使用 FP32 的网络，通常不会损失精度。对于权重或激活需要高动态范围的模型，它比 FP16 更稳健。

有关更多信息，请参考 A100 GPU 中的 TensorFloat-32 将 AI 训练、HPC 加速高达 20 倍博客文章。

TF32 在 NVIDIA Ampere GPU 架构中受支持，并且默认启用。

术语表

Degree (type) - 阶数（类型）

特征的阶数或类型是一个正整数，用于描述当输入在三维空间中旋转时该特征的变换方式。

这与不同旋转阶数的不可约表示相关。

特征的阶数决定了其维度。一个d阶特征的维度为2d+1。

一些常见示例包括：

0阶：1D标量，对旋转具有不变性
1阶：3D向量，根据3D旋转矩阵进行旋转
2阶：5D向量，根据5DWigner-D矩阵进行旋转。这些可以表示对称无迹3x3矩阵。

Fiber - 纤维

纤维可以看作是一组不同类型或阶数（正整数）特征的表示，其中每种特征类型都按照其自身规则进行变换。

在本仓库中，纤维可以视为一个字典，其中键为阶数，值为通道数。

Multiplicity - 重数

给定类型特征的重数是指该特征的通道数量。

Tensor Field Network - 张量场网络

张量场网络是一种等变图卷积，它能够组合不同阶数的特征并生成新的特征，同时借助张量积保持等变性。

Equivariance - 等变性

等变性是模型函数的一种性质，指的是对输入应用对称变换后再计算函数，与先计算函数再对输出应用该变换，得到的结果相同。

在SE(3)-Transformer中，对称群是连续旋转平移群（SE(3)）。

环境设置

以下部分列出了开始训练 SE(3)-Transformer 模型所需满足的条件。

环境要求

本仓库包含一个 Dockerfile，它基于 PyTorch 21.07 NGC 容器进行扩展，并封装了部分依赖项。除了这些依赖项外，请确保您拥有以下组件：

NVIDIA Docker
PyTorch 21.07+ NGC 容器
支持的 GPU：

有关如何开始使用 NGC 容器的更多信息，请参考 NVIDIA GPU Cloud 文档和深度学习文档中的以下部分：

对于无法使用 PyTorch NGC 容器来设置所需环境或创建自己容器的用户，请参考版本化的 NVIDIA 容器支持矩阵。

快速入门指南

克隆仓库。

git clone https://github.com/NVIDIA/DeepLearningExamples
cd DeepLearningExamples/PyTorch/DrugDiscovery/SE3Transformer

构建 se3-transformer PyTorch NGC 容器。
```
docker build -t se3-transformer .
```

在 NGC 容器中启动交互式会话以运行训练/推理。

mkdir -p results
docker run -it --runtime=nvidia --shm-size=8g --ulimit memlock=-1 --ulimit stack=67108864 --rm -v ${PWD}/results:/results se3-transformer:latest

开始训练。
```
bash scripts/train.sh
```
开始推理/预测。
```
bash scripts/predict.sh
```

高级说明

以下部分将详细介绍数据集、训练与推理的运行方式以及训练结果。

脚本与示例代码

在根目录下，最重要的文件包括：

Dockerfile：包含运行 SE(3)-Transformers 所需基本依赖的容器
requirements.txt：运行 SE(3)-Transformers 所需的额外依赖集合
se3_transformer/data_loading/qm9.py：QM9 数据加载、预处理以及基矩阵预计算
se3_transformer/model/layers/：包含模型架构各层的目录
se3_transformer/model/transformer.py：Transformer 主模块
se3_transformer/model/basis.py：基矩阵计算逻辑
se3_transformer/runtime/training.py：训练脚本，需作为 Python 模块运行
se3_transformer/runtime/inference.py：推理脚本，需作为 Python 模块运行
se3_transformer/runtime/metrics.py：支持多 GPU 同步的 MAE 指标
se3_transformer/runtime/loggers.py：DLLogger 和 W&B 日志记录器

参数

training.py 脚本的完整可用参数列表如下：

通用参数

--epochs：训练轮数（默认：单 GPU 下为 100）
--batch_size：批处理大小（默认：240）
--seed：全局随机种子（默认：None）
--num_workers：数据加载工作进程数（默认：8）
--amp：使用自动混合精度（默认：false）
--gradient_clip：梯度范数裁剪（默认：None）
--accumulate_grad_batches：梯度累积（默认：1）
--ckpt_interval：每 N 轮保存一次检查点（默认：-1）
--eval_interval：每 N 轮进行一次评估（默认：1）
--silent：最小化标准输出（默认：false）

路径参数

--data_dir：数据存放或下载目录（默认：./data）
--log_dir：结果日志保存目录（默认：/results）
--save_ckpt_path：检查点保存文件路径（默认：None）
--load_ckpt_path：待加载的检查点文件路径（默认：None）

优化器参数

--optimizer：使用的优化器（默认：adam）
--learning_rate：学习率（默认：单 GPU 下为 0.002）
--momentum：动量（默认：0.9）
--weight_decay：权重衰减（默认：0.1）

QM9 数据集参数

--task：训练的回归任务（默认：homo）
--precompute_bases：在脚本开始时、数据集初始化阶段预计算基矩阵，而非在每次前向传播开始时计算（默认：false）

模型架构参数

--num_layers：堆叠的 Transformer 层数（默认：7）
--num_heads：自注意力头数（默认：8）
--channels_div：送入注意力层前的通道数分割比例（默认：2）
--pooling：图池化类型（默认：max）
--norm：在每个注意力块后应用归一化层（默认：false）
--use_layer_norm：在 MLP 层之间应用层归一化（默认：false）
--low_memory：若为 true，将使用融合操作，速度较慢但内存占用更少（预计减少 25% 内存）。仅在 NVIDIA Volta GPU 上启用 AMP 或在 Ampere GPU 上运行时生效（默认：false）
--num_degrees：使用的度数数量。隐藏特征将具有类型 [0, ..., num_degrees - 1]（默认：4）
--num_channels：隐藏特征的通道数（默认：32）

命令行选项

要查看所有可用选项及其说明的完整列表，请使用 -h 或 --help 命令行选项，例如：python -m se3_transformer.runtime.training --help。

数据集指南

演示数据集

SE(3)-Transformer 是在 QM9 数据集上进行训练的。

QM9 数据集托管在 DGL 服务器上，在需要时会自动下载（38MB）。默认情况下，它存储在 ./data 目录中，但可以通过 --data_dir 参数更改此位置。

该数据集保存为 qm9_edge.npz 文件，并在运行时转换为 DGL 图。

我们使用的输入特征如下：

节点特征（6 维）：
- 原子类型的独热编码（5 维）（原子类型：H、C、N、O、F）
- 每个原子的质子数（1 维）
边特征：键类型的独热编码（4 维）（键类型：单键、双键、三键、芳香键）
相邻节点（原子）之间的相对位置

自定义数据集

要在新数据集上使用此网络，您可以扩展 se3_transformer/data_loading/data_module.py 中提供的 DataModule 类。

您的自定义 collate 函数应返回一个元组，其中包含：

一个（批处理的）DGLGraph 对象
节点特征字典（{'{degree}': tensor}）
边特征字典（{'{degree}': tensor}）
（可选）预计算的基，作为一个字典
标签，作为一个张量

然后，您可以修改 training.py 和 inference.py 脚本以使用您的新数据模块。

训练过程

训练脚本为 se3_transformer/runtime/training.py，需作为模块运行：python -m se3_transformer.runtime.training。

日志

默认情况下，生成的日志存储在 /results/ 中。可以通过 --log_dir 更改此位置。

您可以通过设置 WANDB_API_KEY 环境变量来连接现有的 Weights & Biases 账户。

检查点

可以通过 --save_ckpt_path 参数设置保存检查点的文件路径。还可以通过 --ckpt_interval 设置检查点之间的间隔（以 epoch 数为单位）。

评估

评估指标为平均绝对误差（MAE）。

--eval_interval 可以设置评估轮次之间的间隔（以 epoch 数为单位）。默认情况下，每个 epoch 后执行一次评估。

自动混合精度

要启用混合精度训练，请添加 --amp 标志。

多 GPU 和多节点

训练脚本支持 PyTorch elastic 启动器，以在多个 GPU 或节点上运行。请参考官方文档。

例如，要使用 AMP 在所有可用 GPU 上进行训练：

python -m torch.distributed.run --nnodes=1 --nproc_per_node=gpu --module se3_transformer.runtime.training --amp

推理流程

可通过 se3_transformer.runtime.inference Python 模块运行推理。

性能

基准测试

以下部分介绍如何运行基准测试以衡量模型在训练和推理模式下的性能。

训练性能基准测试

推理性能基准测试

要在特定批处理大小下对推理性能进行基准测试，运行 bash scripts/benchmarck_inference.sh {BATCH_SIZE}。

结果

以下部分详细介绍我们在训练和推理中如何实现性能和精度。

训练精度结果

训练精度：NVIDIA DGX A100（8x A100 80GB）

我们的结果是通过在配备 8x A100 80GB GPU 的 NVIDIA DGX A100 上，在 PyTorch 21.07 NGC 容器中运行 scripts/train.sh 训练脚本获得的。

GPU 数量	每 GPU 批处理大小	TF32 绝对误差	混合精度绝对误差	TF32 训练时间	混合精度训练时间	混合精度相对 TF32 的训练速度提升
1	240	0.03456	0.03460	1小时23分钟	1小时03分钟	1.32倍
8	240	0.03417	0.03424	15分钟	12分钟	1.25倍

训练精度：NVIDIA DGX-1（8x V100 16GB）

我们的结果是通过在配备 8x V100 16GB GPU 的 NVIDIA DGX-1 上，在 PyTorch 21.07 NGC 容器中运行 scripts/train.sh 训练脚本获得的。

GPU 数量	每 GPU 批处理大小	FP32 绝对误差	混合精度绝对误差	FP32 训练时间	混合精度训练时间	混合精度相对 FP32 的训练速度提升
1	240	0.03432	0.03439	2小时25分钟	1小时33分钟	1.56倍
8	240	0.03380	0.03495	29分钟	20分钟	1.45倍

训练性能结果

训练性能：NVIDIA DGX A100（8x A100 80GB）

GPU 数量	每 GPU 批处理大小	TF32 吞吐量 [mol/ms]	混合精度吞吐量 [mol/ms]	混合精度相对 TF32 的吞吐量提升	TF32 弱扩展性	混合精度弱扩展性
1	240	2.21	2.92	1.32倍
1	120	1.81	2.04	1.13倍
8	240	17.15	22.95	1.34倍	7.76	7.86
8	120	13.89	15.62	1.12倍	7.67	7.66

要获得相同结果，请按照快速入门指南中的步骤操作。

训练性能：NVIDIA DGX-1（8x V100 16GB）

GPU 数量	每 GPU 批处理大小	FP32 吞吐量 [mol/ms]	混合精度吞吐量 [mol/ms]	混合精度相对 FP32 的吞吐量提升	FP32 弱扩展性	混合精度弱扩展性
1	240	1.25	1.88	1.50倍
1	120	1.03	1.41	1.37倍
8	240	9.33	14.02	1.50倍	7.46	7.46
8	120	7.39	9.41	1.27倍	7.17	6.67

要获得相同结果，请按照快速入门指南中的步骤操作。

推理性能结果

推理性能：NVIDIA DGX A100（1x A100 80GB）

我们的结果是通过在配备 1x A100 80GB GPU 的 NVIDIA DGX A100 上，在 PyTorch 21.07 NGC 容器中运行 scripts/benchmark_inference.sh 推理基准测试脚本获得的。

FP16

批处理大小	平均吞吐量 [mol/ms]	平均延迟 [ms]	90% 延迟 [ms]	95% 延迟 [ms]	99% 延迟 [ms]
1600	11.60	140.94	138.29	140.12	386.40
800	10.74	75.69	75.74	76.50	79.77
400	8.86	45.57	46.11	46.60	49.97

TF32

批处理大小	平均吞吐量 [mol/ms]	平均延迟 [ms]	90% 延迟 [ms]	95% 延迟 [ms]	99% 延迟 [ms]
1600	8.58	189.20	186.39	187.71	420.28
800	8.28	97.56	97.20	97.73	101.13
400	7.55	53.38	53.72	54.48	56.62

要获得相同结果，请按照快速入门指南中的步骤操作。

推理性能：NVIDIA DGX-1（1x V100 16GB）

我们的结果是通过在配备 1x V100 16GB GPU 的 NVIDIA DGX-1 上，在 PyTorch 21.07 NGC 容器中运行 scripts/benchmark_inference.sh 推理基准测试脚本获得的。

FP16

批处理大小	平均吞吐量 [mol/ms]	平均延迟 [ms]	90% 延迟 [ms]	95% 延迟 [ms]	99% 延迟 [ms]
1600	6.42	254.54	247.97	249.29	721.15
800	6.13	132.07	131.90	132.70	140.15
400	5.37	75.12	76.01	76.66	79.90

FP32

批处理大小	平均吞吐量 [mol/ms]	平均延迟 [ms]	90% 延迟 [ms]	95% 延迟 [ms]	99% 延迟 [ms]
1600	3.39	475.86	473.82	475.64	891.18
800	3.36	239.17	240.64	241.65	243.70
400	3.17	126.67	128.19	128.82	130.54

要获得相同结果，请按照快速入门指南中的步骤操作。

发布说明

更新日志

2021年8月

初始版本发布