NVIDIA cuEquivariance
cuEquivariance 是一款 CUDA-X™ 库，专为应对几何感知神经网络在 3D 数据任务中所需的高计算需求而设计。该库提供了经过优化的 NVIDIA CUDA® 内核和完善的 API，包括三角注意力（triangle attention）和三角乘法（triangle multiplication）等功能，以加速这些过程在药物和材料发现等各类科学领域中的应用。
主要特性
灵活的 API
其他等变性库通常受限于对 SO(3) 不可约表示（irreps）基底和数据布局的特定选择。而使用 cuEquivariance，您可以通过创建分段张量积（segmented tensor product）自定义不可约表示基底的张量积，并将这类操作推广到不可约表示之外，从而构建等变神经网络。
CUDA 加速性能
实现以下目标：
- 10x speedup for end-to-end MACE performance
- 200x speedup for symmetric contraction operation performance
- 100,000 natoms per GPU being simulated with MACE
- 3.5x speedups for triangle operations performance
有关上述性能的更多信息，请查看下面的“性能”部分。
扩展 MLIP 支持和加速
主流的等变性机器学习原子间势能模型，包括 MACE、Allegro、NequIP 和 DiffDock
包含三角核的蛋白质模型，包括：Boltz、Neo-1 和 OpenFold
开始使用 NVIDIA cuEquivariance
使用 Conda 快速安装
conda install conda-forge::cuequivariance
使用 pip 快速安装
# Choose the frontend you want to use pip install cuequivariance-jax pip install cuequivariance-torch pip install cuequivariance # Installs only the core non-ML components # CUDA kernels pip install cuequivariance-ops-torch-cu11 pip install cuequivariance-ops-torch-cu12 pip install cuequivariance-ops-jax-cu12
性能
更多资源
伦理 AI
NVIDIA 认为可信 AI 是一项共同责任，我们已制定相关政策和实践，以支持各种 AI 应用的开发。根据我们的服务条款下载或使用时，开发者应与其支持的模型团队合作，确保此模型满足相关行业和用例的要求，并解决不可预见的产品滥用问题。
如需了解有关本模型伦理考量的详细信息，请参阅 Model Card++ 的可解释性、偏见、安全与隐私等子卡片。如需报告安全漏洞或 NVIDIA AI 相关问题，请点击此处。
立即开始使用 cuEquivariance。