Skip to main content

A pytorch package for robustness and adaptation research

Project description

PWC PWC PWC

Robustness evaluation and Adaptation of ImageNet models

This repo contains a growing collection of helper functions, tools and methods for robustness evaluation and adaptation of ImageNet scale models. The focus is on simple methods that work at scale.

We currently have the following features available:

  • examples/batchnorm: A reference implementation of batch norm adaptation used by Schneider, Rusak et al. (NeurIPS 2020)
  • examples/selflearning: A reference implementation of self learning with robust pseudo labeling used by Rusak, Schneider et al. (arxiv 2021)
  • examples/imagenet_d: Example runs on the ImageNet-D dataset used by Rusak, Schneider et al. (arxiv 2021)

Planned features for future releases are (please open an issue if you can think of additional interesting parts to add):

  • Helper functions for robustness datasets like ImageNet-A, ImageNet-R and ImageNet-C
  • examples/clip: Robustness evaluation for CLIP, Radford et al. (2021)
  • examples/dino: Robustness evaluation for DINO, Caron et al. (2021)

News

  • May '21: We will present our work on self-learing as a contributed talk at the WeaSuL 2021 workshop at ICLR.
  • April '21: The pre-print for "Adapting ImageNet-scale models to complex distribution shifts with self-learning" is now available on arXiv: arxiv.org/abs/2104.12928
  • September 2020: The BatchNorm adaptation paper was accepted for poster presentation at NeurIPS 2020.
  • July '20: A shorter workshop version of the paper was accepted for oral presentation at the Uncertainty & Robustness in Deep Learning Workshop at ICML 2020.
  • June '20: The pre-print for "Improving robustness against common corruptions by covariate shift adaptation" is available on arXiv: arxiv.org/abs/2006.16971.pdf

☕ The robusta toolbox for Robustness and Adaptation

Motivation

Besides reference implementations, this repo is mainly intended to provide a quick and easy way to adapt your own code. In particular, when developing new methods for improving robustness on deep learning models, we find it interesting to report results after adapting your model to the test datasets. This paints a more holistic image of model robustness: Some people might be interested in ad-hoc model performance, other might be interested in the performance obtained in a transductive inference setting.

Note that the package is not intended for general purpose domain adaptation. Instead, we focus on providing simple methods that prove to be effective for ImageNet scale model adaptation at test time. The package provides helper functions that are "minimally invasive" and can easily be added to existing source code for model evaluation.

Quick Start

We will release the first stable version of the package on PyPI. Until then, you can install directly from the main repo:

pip install git+git://github.com/bethgelab/robustness.git

Here is an example for how to use robusta for batchnorm adaptation & robust pseudo-labeling.

    model = torchvision.models.resnet50(pretrained=True)

    # We provide implementations for ImageNet-val, ImageNetC, ImageNetR,
    # ImageNetA and ImageNetD:
    val_dataset = robusta.datasets.imagenetc.ImageNetC(
        root=dataset_folder, corruption="gaussian_blur", severity=1,
        transform=transforms.Compose([transforms.ToTensor()])
        )
    val_loader = torch.utils.data.DataLoader(
        val_dataset, batch_size=batch_size, shuffle=True)

    # We offer different options for batch norm adaptation;
    # alternatives are "ema", "batch_wise_prior", ...
    robusta.batchnorm.adapt(model, adapt_type="batch_wise")

    # The accuracy metric can be specific to the dataset:
    # For example, ImageNet-R requires remapping into 200 classes.
    accuracy_metric = val_dataset.accuracy

    # You can also easily use self-learning in your model.
    # Self-learning adaptation can be combined with batch norm adaptation, example:
    parameters = robusta.selflearning.adapt(model, adapt_type="affine")
    optimizer = torch.optim.SGD(parameters, lr=1e-3)

    # You can choose from a set of adaptation losses (GCE, Entropy, ...)
    rpl_loss = robusta.selflearning.GeneralizedCrossEntropy(q=0.8)

    acc1_sum, acc5_sum, num_samples = 0., 0., 0.
    for epoch in range(num_epochs):
        predictions = []
        for images, targets in val_loader:

            logits = model(images)
            predictions = logits.argmax(dim=1)

            # Predictions are optional. If you do not specify them,
            # they will be computed within the loss function.
            loss = rpl_loss(logits, predictions)

            # When using self-learning, you need to add an additional optimizer
            # step in your evaluation loop.
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            acc1_sum, acc5_sum += accuracy_metric(predictions, targets, topk=(1,5))
            num_samples += len(targets)
            print(f"Top-1: {acc1_sum/num_samples}, Top-5: {acc5_sum/num_samples}")

Example Implementations

Batch Norm Adaptation

PWC PWC

[Paper] [Web] [README] [Implementation]

We propose to go beyond the assumption of a single sample from the target domain when evaluating robustness. Re-computing BatchNorm statistics is a simple baseline algorithm for improving the corruption error up to 14% points over a wide range of models, when access to more than a single sample is possible.

Self-Learning

PWC PWC PWC

[Paper] [Web] [README] [Implementation]

Test-time adaptation with self-learning improves robustness of large-scale computer vision models on ImageNet-C, -R, and -A.

Robustness evaluation of DINO models

[Blog Post] [Implementation coming soon]

License

Unless noted otherwise, code in this repo is released under an Apache 2.0 license. Some parts of the implementation use third party code. We typically indicate this in the file header or in the methods directly, and include the original license in the NOTICE file.

This repo does not contain the full code-base used in Rusak, Schneider et al. (2021) and is instead currently limited to a reference re-implementation for robust-pseudo labeling and entropy minimization. A full version of the codebase might be independently released in the future.

If you want to use part of this code commercially, please carefully check the involved parts. Part of the third-party implementations might be released under licences with a non-commercial use clause such as CC-NC. If in doubt, please reach out.

Contact

Please reach out for feature requests. Contributions welcome!

Note: The current version of this code base is a work in progress. We still decided to do this pre-release since the core methods are conceptually easy to use in your own code (batch norm adaptation, self-learning, ... and the current state might already be a useful place to start.

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

robusta-0.0.2.tar.gz (34.9 kB view details)

Uploaded Source

Built Distribution

robusta-0.0.2-py3-none-any.whl (43.6 kB view details)

Uploaded Python 3

File details

Details for the file robusta-0.0.2.tar.gz.

File metadata

  • Download URL: robusta-0.0.2.tar.gz
  • Upload date:
  • Size: 34.9 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/4.0.2 CPython/3.8.17

File hashes

Hashes for robusta-0.0.2.tar.gz
Algorithm Hash digest
SHA256 5000d7ba38da148d304910e92f5f971bd1ce103db3d7748e6a61f7906af6cbde
MD5 b3a75cfafd72bf356d6a1d7895142ddd
BLAKE2b-256 d37e7d4007b7ae961bdf05640330e207979d6c258115aeb358529ff3eb0a6398

See more details on using hashes here.

File details

Details for the file robusta-0.0.2-py3-none-any.whl.

File metadata

  • Download URL: robusta-0.0.2-py3-none-any.whl
  • Upload date:
  • Size: 43.6 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/4.0.2 CPython/3.8.17

File hashes

Hashes for robusta-0.0.2-py3-none-any.whl
Algorithm Hash digest
SHA256 aa324cbfbb7592fb82a13e7d21e0ee9cde0f69edf7bc41ee8204a0ee50a2aa3f
MD5 aeaabe1667ff7d5155c2d536f3db2110
BLAKE2b-256 a8fc104e453ec14ac0f7e8b7673ab192c7466cb0e5f40f2daeedb0cfbae79003

See more details on using hashes here.

Supported by

AWS AWS Cloud computing and Security Sponsor Datadog Datadog Monitoring Fastly Fastly CDN Google Google Download Analytics Microsoft Microsoft PSF Sponsor Pingdom Pingdom Monitoring Sentry Sentry Error logging StatusPage StatusPage Status page