Create environment
Create conda environment:
conda create --name clatr python=3.10 -y
conda activate clatr
Install dependencies:
conda install pytorch pytorch-cuda=12.1 -c pytorch -c nvidia
pip install -r requirements.txt
Set up the dataset
Check the data repository here.
First, install git lfs by following the instructions here.
To get the data, run:
cd /PATH/TO/THE/DATASET
git clone https://huggingface.co/datasets/robin-courant/et-data
Prepare the dataset (untar archives):
cd et-data
sh untar_and_move.sh
Then copy or symlink the et-data folder in ./data:
cd /PATH/TO/CLaTr
ln -s /PATH/TO/THE/DATASET/et-data ./data
Download checkpoints
Run:
sh download_checkpoints.sh
Launch training:
HYDRA_FULL_ERROR=1 python src/train.py
Launch evaluation:
HYDRA_FULL_ERROR=1 python src/evaluate.py checkpoint_path=checkpoints/clatr-e100.ckpt
Launch feature extraction:
HYDRA_FULL_ERROR=1 python src/extraction.py checkpoint_path=checkpoints/clatr-e100.ckpt
Extract checkpoint into submodule checkpoints (traj_encoder, text_encoder, traj_decoder):
HYDRA_FULL_ERROR=1 python src/extract_ckpt.py checkpoint_path=checkpoints/clatr-e100.ckpt
This code is distributed under an MIT LICENSE.
Note that our code depends on other libraries, including PyTorch, PyTorch3D, Hugging Face, Hydra, and uses datasets which each have their own respective licenses that must also be followed.
Note also that this repository is highly adapted from the codebase of TMR by Mathis Petrovich.