$\texttt{skwdro}$: a library for Wasserstein distributionally robust machine learning

TL;DR

skwdro is a Python library that simplifies training robust machine learning models using Wasserstein distributionally robust optimization, integrating with PyTorch and scikit-learn for ease of use.

Contribution

The paper introduces skwdro, a library that makes Wasserstein distributionally robust optimization accessible with minimal code modifications.

Findings

Provides a user-friendly wrapper for robust training in PyTorch

Includes scikit-learn compatible estimators for common objectives

Uses entropic smoothing to enhance model flexibility

Abstract

We present skwdro, a Python library for training robust machine learning models. The library is based on distributionally robust optimization using Wasserstein distances, popular in optimal transport and machine learnings. The goal of the library is to make the training of robust models easier for a wide audience by proposing a wrapper for PyTorch modules, enabling model loss' robustification with minimal code changes. It comes along with scikit-learn compatible estimators for some popular objectives. The core of the implementation relies on an entropic smoothing of the original robust objective, in order to ensure maximal model flexibility. The library is available at https://github.com/iutzeler/skwdro and the documentation at https://skwdro.readthedocs.io.