SE(3)-Transformers for PyTorch

NVIDIA

Resource

NVIDIA

SE(3)-Transformers for PyTorch

A Graph Neural Network using a variant of self-attention for 3D points and graphs processing.

To train your model using mixed or TF32 precision with Tensor Cores or FP32, perform the following steps using the default parameters of the SE(3)-Transformer model on the QM9 dataset. For the specifics concerning training and inference, refer to the Advanced section.

Clone the repository.

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

Build the se3-transformer PyTorch NGC container.
```
docker build -t se3-transformer .
```

Start an interactive session in the NGC container to run training/inference.

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

Start training.
```
bash scripts/train.sh
```
Start inference/predictions.
```
bash scripts/predict.sh
```

Now that you have your model trained and evaluated, you can choose to compare your training results with our Training accuracy results. You can also choose to benchmark your performance to Training performance benchmark or Inference performance benchmark. Following the steps in these sections will ensure that you achieve the same accuracy and performance results as stated in the Results section.