This repository contains a set of tools for training models with the SHIFT dataset. The SHIFT dataset is a multi-task synthetic driving dataset that was created using the CARLA simulator. As such, it provides an opportunity to train models which can be deployed in autonomous driving stacks operating inside CARLA, providing a rich research opportunity in AV stack development.
-
Multiformer - read the blog post and WandB Report
- Expansion of Multitask Segformer into 2D object detection by incorporating a scaled-down version of Deformable DETR.
- Usage scripts:
-
Multitask Segformer - read the blog post and WandB Report
- Uses a PVTv2 backbone with two task heads (possibly could be expanded to include others). Selection of B0-B5 model sizes for desired accuracy/efficiency tradeoff.
- Usage scripts:
If you'd prefer to work outside the Docker container, you can set up like this:
- Install latest NVIDIA driver for your hardware.
- Install CUDA Toolkit.
- Set up a virtual environment:
python -m venv ./venv
- Activate virtual environment:
- Windows Powershell:
./venv/Scripts/Activate.ps1
- Linux:
source ./venv/bin/activate
- Windows Powershell:
- Install repo:
- Set environment variable
- Linux:
export PIP_EXTRA_INDEX_URL=https://download.pytorch.org/whl/cu117
- Windows:
$env:PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/cu117"
- Linux:
pip install -e .
- Set environment variable
- Install bitsandbytes (optional):
- git clone https://github.com/TimDettmers/bitsandbytes.git
- cd bitsandbytes
- CUDA_VERSION=117 make cuda11x
- python setup.py install
- cd ..
- Run a training script (Use --help to discover params):
python ./scripts/model_train_eval/train_multiformer.py --help
Navigate in your terminal into this repository. With Docker installed on your system:
- Build Docker image (this may take a while):
docker build -f Dockerfile . -t shift_lab:latest
- Run Docker image:
docker run -it --gpus all --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 shift_lab:latest
- Run a training script (Use --help to discover params):
python shift-experiments/scripts/model_train_eval/train_multitask_segformer.py --help