CUDA Equivariance for PyTorch

0.9.1 · active · verified Thu Apr 16

cuequivariance-torch provides CUDA-accelerated implementations of equivariant operations for PyTorch. It aims to efficiently handle geometric symmetries in deep learning models, particularly for 3D data, by offering modules for SO(3) rotations and SE(3) translations. The current version is 0.9.1, and it maintains an active development pace with new features and optimizations.

Common errors

```
RuntimeError: CUDA is not available
```
cause PyTorch was installed without CUDA support, or there is no CUDA-enabled GPU detected on the system.

fix Reinstall PyTorch with CUDA support (check official PyTorch installation instructions for your CUDA version) and ensure your system has a properly configured CUDA-enabled GPU with up-to-date drivers.
```
ModuleNotFoundError: No module named 'cuequivariance_torch'
```
cause Attempting to import the library using its PyPI package name (`cuequivariance-torch`) rather than its internal module name (`cuequivariance`).

fix Change your import statements from `import cuequivariance_torch` or `from cuequivariance_torch import ...` to `import cuequivariance` or `from cuequivariance import ...`.
```
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu
```
cause One or more input tensors or the neural network module itself have not been moved to the CUDA device before computation.

fix Ensure all tensors involved in an operation are on the same device. Move your input tensors using `.to(device)` (e.g., `input_data.to('cuda')`) and your model using `.to(device)` (e.g., `model.to('cuda')`). Define `device = torch.device('cuda')` and consistently apply it.

Warnings

breaking The library explicitly requires PyTorch 2.1 or newer. Older versions of PyTorch or PyTorch installations without CUDA will not work.
Fix: Ensure your PyTorch installation is `torch>=2.1` and is built with CUDA support (e.g., install `torch` using a command that specifies CUDA, such as from the official PyTorch website).
gotcha The PyPI package name is `cuequivariance-torch`, but the Python import name is `cuequivariance`.
Fix: Always use `from cuequivariance import ...` or `from cuequivariance.modules import ...` for imports, not `cuequivariance_torch`.
gotcha This library is CUDA-accelerated and strictly requires a CUDA-enabled GPU and a PyTorch installation compiled with CUDA. It will not function on CPU-only setups.
Fix: Verify that `torch.cuda.is_available()` returns `True`. If not, ensure you have a CUDA-compatible GPU, the necessary drivers, and PyTorch installed with CUDA support.

Install

pip install cuequivariance-torch Install stable version

Imports

SO3Linear

from cuequivariance.modules import SO3Linear

SO3RotMat
wrong:
```
from cuequivariance_torch import SO3RotMat
```
correct:
```
from cuequivariance import SO3RotMat
```
The top-level package for imports is `cuequivariance`, not `cuequivariance-torch` (the PyPI package name).

Quickstart

This quickstart demonstrates how to initialize and use an `SO3Linear` layer, a core module for SO(3) equivariant operations, ensuring that both the input tensor and the model are correctly placed on a CUDA device.

import torch
from cuequivariance.modules import SO3Linear

# Ensure CUDA is available
if not torch.cuda.is_available():
    raise RuntimeError("CUDA is not available. cu-equivariance-torch requires a CUDA-enabled PyTorch installation with a GPU.")

# Define device
device = torch.device("cuda")

# Example: SO(3) Linear layer for rotation matrices
# Input dimensions: (batch, features_in, 3, 3) where the last two dimensions represent a 3x3 matrix
batch_size = 4
features_in = 16
features_out = 32

# Create a dummy input tensor on the CUDA device
# For SO3Linear, the input tensor's last two dimensions are treated as matrix components.
# The layer handles the equivariant operations.
input_data = torch.randn(batch_size, features_in, 3, 3, device=device)

# Instantiate an SO(3) Linear layer and move it to the CUDA device
so3_linear_layer = SO3Linear(features_in, features_out).to(device)

# Pass the input through the layer
output_data = so3_linear_layer(input_data)

print(f"Input shape: {input_data.shape}, device: {input_data.device}")
print(f"Output shape: {output_data.shape}, device: {output_data.device}")
# Expected output shape: (batch_size, features_out, 3, 3)

view raw JSON →