# SE(3)-Equivariant Attention Networks for Shape Reconstruction in Function Space
This a preliminary release of the code for the paper "SE(3)-Equivariant Attention Networks for Shape Reconstruction in Function Space" 

The code structure was based on the following repositories:
* [DeepLearningExamples/DGLPyTorch/DrugDiscovery/SE3Transformer](https://github.com/NVIDIA/DeepLearningExamples/tree/master/DGLPyTorch/DrugDiscovery/SE3Transformer)
* [ autonomousvision/occupancy_networks](https://github.com/autonomousvision/occupancy_networks)
# Data Preparation
Download the ShapeNet dataset using the instructions found in the [Occupancy Networks](https://github.com/autonomousvision/occupancy_networks) repository.

If you want to train or evaluate on a rotated version of the dataset you need to run the script to generate a random rotation for each shape
```
python scripts/generate_random_rotation.py 
```

# Run on AWS Sagemaker
You need to upload the ShapeNet dataset in the 's3://\<s3-backet-uri\>/data', where \<s3-backet-uri\> is your backet's uri.

If you have  generated the random rotations using the script [generate_random_rotation.py](script/generate_random_rotation.py) you must upload them  in the 's3://\<s3-backet-uri\>/checkpoints/rotations'

## Training and Inference on AWS
We provide the python notebook [training_sagemaker.ipynb](training_sagemaker.ipynb) with the default settings for training on AWS. To perfom the training with the Pytorch estimator:
```
from sagemaker.pytorch import PyTorch

experiment_name='DefaultModel'
backet_uri='<s3-backet-uri>'
estimator = PyTorch(
    entry_point="recon_inference.py",
    role=role,
    framework_version="1.9.0",
    instance_count=1,
    instance_type="ml.p3.16xlarge",
    source_dir=".",
    py_version="py38",
    hyperparameters=parameters,
    input_mode='FastFile',
    distribution={"smdistributed": {"dataparallel": {"enabled": True}}},
    
    checkpoint_s3_uri='s3://'+backet_uri+'/checkpoints/'+experiment_name,
    checkpoint_local_path='/opt/ml/checkpoints'
)
```

You need to replace \<s3-backet-uri\> with your s3 backet.

Similarly you can perform inference using the [inference_sagemaker.ipynb](inference_sagemaker.ipynb) python notebook
  
# Run Locally
This code was partially tested for local execution. We provide some general installation instructions and running scripts. For more detailed installation instructions please refer [here](https://github.com/NVIDIA/DeepLearningExamples/tree/master/DGLPyTorch/DrugDiscovery/SE3Transformer#quick-start-guide)
## Installation 
  1. Build the se3-transformer PyTorch NGC container.
  ```
  docker build -t se3-transformer .
  ```
  2. 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
  ```
 ## Training and Inference Locally
  1. Train using the following script
  ```
  bash scripts/train.sh
  ```
  2. Perform inference using the following script 
  ```
  bash scripts/predict.sh
  ```
