Overview

WILDS is a benchmark of in-the-wild distribution shifts spanning diverse data modalities and applications, from tumor identification to wildlife monitoring to poverty mapping.

The WILDS package contains:

Data loaders that automatically handle data downloading, processing, and splitting, and
Dataset evaluators that standardize model evaluation for each dataset.

In addition, the example scripts contain default models, optimizers, schedulers, and training/evaluation code. New algorithms can be easily added and run on all of the WILDS datasets.

For more information, please visit our website or read the main WILDS paper (1) and its follow-up integrating unlabeled data (2). For questions and feedback, please post on the discussion board.

Installation

We recommend using pip to install WILDS:

pip install wilds

If you have already installed it, please check that you have the latest version:

python -c "import wilds; print(wilds.__version__)"
# This should print "2.0.0". If it doesn't, update by running:
pip install -U wilds

If you plan to edit or contribute to WILDS, you should install from source:

git clone git@github.com:p-lambda/wilds.git
cd wilds
pip install -e .

In examples/, we provide a set of scripts that can be used to train models on the WILDS datasets. These scripts were also used to benchmark baselines in our papers [1, 2]. These scripts are not part of the installed WILDS package. To use them, you should install from source, as described above.

Requirements

The WILDS package depends on the following requirements:

numpy>=1.19.1
ogb>=1.2.6
outdated>=0.2.0
pandas>=1.1.0
pillow>=7.2.0
pytz>=2020.4
torch>=1.7.0
torch-scatter>=2.0.5
torch-geometric>=2.0.1
torchvision>=0.8.2
tqdm>=4.53.0
scikit-learn>=0.20.0
scipy>=1.5.4

Running pip install wilds or pip install -e . will automatically check for and install all of these requirements except for the torch-scatter and torch-geometric packages, which require a quick manual install.

Example script requirements

To run the example scripts, you will also need to install these additional dependencies:

transformers>=3.5.0
SwAV requires Apex. To install Apex, please follow the README from the official SwAV repository.
Our code supports the optional use of Weights & Biases to track and monitor experiments. To install the Weights and Biases Python package, run pip install wandb.

All baseline experiments in the paper were run on Python 3.8.5 and CUDA 10.1.

Datasets

WILDS currently includes 10 datasets, which we've briefly listed below. For full dataset descriptions, please see our papers (1, 2).

| Dataset | Modality | Labeled splits | Unlabeled splits | | ------------- | -------- | --------------------------------- | --------------------------------------------------------------- | | iwildcam | Image | train, val, test, id_val, id_test | extra_unlabeled | | camelyon17 | Image | train, val, test, id_val | train_unlabeled, val_unlabeled, test_unlabeled | | rxrx1 | Image | train, val, test, id_test | - | | ogb-molpcba | Graph | train, val, test | train_unlabeled, val_unlabeled, test_unlabeled | | globalwheat | Image | train, val, test, id_val, id_test | train_unlabeled, val_unlabeled, test_unlabeled, extra_unlabeled | | civilcomments | Text | train, val, test | extra_unlabeled | | fmow | Image | train, val, test, id_val, id_test | train_unlabeled, val_unlabeled, test_unlabeled | | poverty | Image | train, val, test, id_val, id_test | train_unlabeled, val_unlabeled, test_unlabeled | | amazon | Text | train, val, test, id_val, id_test | val_unlabeled, test_unlabeled, extra_unlabeled | | py150 | Text | train, val, test, id_val, id_test | - |

Using the WILDS package

Data

The WILDS package provides a simple, standardized interface for all datasets in the benchmark. This short Python snippet covers all of the steps of getting started with a WILDS dataset, including dataset download and initialization, accessing various splits, and preparing a user-customizable data loader. We discuss data loading in more detail in #Data loading.

from wilds import get_dataset
from wilds.common.data_loaders import get_train_loader
import torchvision.transforms as transforms

# Load the full dataset, and download it if necessary
dataset = get_dataset(dataset="iwildcam", download=True)

# Get the training set
train_data = dataset.get_subset(
    "train",
    transform=transforms.Compose(
        [transforms.Resize((448, 448)), transforms.ToTensor()]
    ),
)

# Prepare the standard data loader
train_loader = get_train_loader("standard", train_data, batch_size=16)

# (Optional) Load unlabeled data
dataset = get_dataset(dataset="iwildcam", download=True, unlabeled=True)
unlabeled_data = dataset.get_subset(
    "test_unlabeled",
    transform=transforms.Compose(
        [transforms.Resize((448, 448)), transforms.ToTensor()]
    ),
)
unlabeled_loader = get_train_loader("standard", unlabeled_data, batch_size=16)

# Train loop
for labeled_batch, unlabeled_batch in zip(train_loader, unlabeled_loader):
    x, y, metadata = labeled_batch
    unlabeled_x, unlabeled_metadata = unlabeled_batch
    ...

The metadata contains information like the domain identity, e.g., which camera a photo was taken from, or which hospital the patient's data came from, etc., as well as other metadata.

Domain information

To allow algorithms to leverage domain annotations as well as other groupings over the available metadata, the WILDS package provides Grouper objects. These Grouper objects are helper objects that extract group annotations from metadata, allowing users to specify the grouping scheme in a flexible fashion. They are used to initialize group-aware data loaders (as discussed in #Data loading) and to implement algorithms that rely on domain annotations (e.g., Group DRO). In the following code snippet, we initialize and use a Grouper that extracts the domain annotations on the iWildCam dataset, where the domain is location.

from wilds.common.grouper import CombinatorialGrouper

# Initialize grouper, which extracts domain information
# In this example, we form domains based on location
grouper = CombinatorialGrouper(dataset, ['location'])

# Train loop
for x, y_true, metadata in train_loader:
    z = grouper.metadata_to_group(metadata)
    ...

Data loading

For training, the WILDS package provides two types of data loaders. The standard data loader shuffles examples in the training set, and is used for the standard approach of empirical risk minimization (ERM), where we minimize the average loss.

from wilds.common.data_loaders import get_train_loader

# Prepare the standard data loader
train_loader = get_train_loader('standard', train_data, batch_size=16)

To support other algorithms that rely on specific data loading schemes, we also provide the group data loader. In each minibatch, the group loader first samples a specified number of groups, and then samples a fixed number of examples from each of those groups. (By default, the groups are sampled uniformly at random, which upweights minority groups as a result. This can be toggled by the uniform_over_groups parameter.) We initialize group loaders as follows, using Grouper that specifies the grouping scheme.

# Prepare a group data loader that samples from user-specified groups
train_loader = get_train_loader(
    "group", train_data, grouper=grouper, n_groups_per_batch=2, batch_size=16
)

Lastly, we also provide a data loader for evaluation, which loads examples without shuffling (unlike the training loaders).

from wilds.common.data_loaders import get_eval_loader

# Get the test set
test_data = dataset.get_subset(
    "test",
    transform=transforms.Compose(
        [transforms.Resize((224, 224)), transforms.ToTensor()]
    ),
)

# Prepare the evaluation data loader
test_loader = get_eval_loader("standard", test_data, batch_size=16)

Evaluators

The WILDS package standardizes and automates evaluation for each dataset. Invoking the eval method of each dataset yields all metrics reported in the paper and on the leaderboard.

from wilds.common.data_loaders import get_eval_loader

# Get the test set
test_data = dataset.get_subset(
    "test",
    transform=transforms.Compose(
        [transforms.Resize((224, 224)), transforms.ToTensor()]
    ),
)

# Prepare the data loader
test_loader = get_eval_loader("standard", test_data, batch_size=16

Wilds

Install / Use

README