Skip to content

Risk model training code for Covid-19 tracing application.


Notifications You must be signed in to change notification settings


Repository files navigation


This repository provides models, infrastructure and datasets for training deep-learning based predictors of COVID-19 infectiousness as used in Proactive Contact Tracing.

What's in the Box


We provide architectural scaffolding around deep-sets (DS), set-transformers (ST) and DS-ST hybrids.

Training Infrastructure

The core training infrastructure is built with speedrun. It supports experiment tracking with Weights and Biases & tensorboard, and hyperparameter sweeps with Weights and Biases Sweeps.


The training data is derived from COVI-AgentSim, an agent based simulator built with contact-tracing benchmarking and epidemiological realism in mind. We will include a dataset (download to appear soon); but in the mean time, you can print your own datasets by following instructions in the COVI-AgentSim repository.

The training and validation data is structured in directories containing zarr datasets. There is some flexibility here (see below), but we expect two directories ./data/train and ./data/val, where each directory should contain an arbitrary number of zarr files.


This repository works with PyTorch. Install the dependencies with:

pip install -r requirements.txt

Please make sure you have a Weights and Biases account and it is configured correctly. If you do not wish to use Weights and Biases, you will need to run the following command before running the main training script:

export WANDB_MODE=dryrun


To run your training script, first make a directory where your experiment logs will live. Once you're in this repository,

mkdir experiments

If your data lives in ./data, you may use the following command to train DS-PCT:

python experiments/DS-PCT-0 --inherit base_configs/DS-PCT-X

... or the following to run ST-PCT:

python experiments/ST-PCT-0 --inherit base_configs/ST-PCT-X

If your data lives elsewhere, you will need to run the following command:

python experiments/DS-PCT-0 --inherit base_configs/DS-PCT-X path/to/training/data path/to/validation/data

Visualizing results

If you are using Weights and Biases, you should have a project named ctt under your account. Additionally, tensorboard logs are dumped in experiments/DS-PCT-0/Logs and checkpoints are stored in experiments/DS-PCT-0/Weights (likewise for ST-PCT).

Reporting bugs and getting help

If you find a bug or have a question, please open an issue in this repository.


Risk model training code for Covid-19 tracing application.







No releases published


No packages published