art-resampling is a PyTorch utility for Adaptive Resampling-Based Training for Imbalanced Datasets (ART).
This repository contains the reference implementation used in the paper: Adaptive Resampling for Imbalanced Datasets (ART)
pip install art-resampling
Run the example script:
python examples/tabular_multiclass.py
A PyTorch Dataset that can act as:
- a normal dataset (
enable_art=False), or - an ART-resampled dataset (
enable_art=True).
from art_resampling import ARTDataset train_dataset = ARTDataset( X_train, y_train, enable_art=True, c=0.3, cls_weights=initial_class_weights, seed=5529, ) val_dataset = ARTDataset(X_val, y_val, enable_art=False)
from art_resampling import art_refresh_dataset_from_predictions train_loader, class_weights, per_class_f1 = art_refresh_dataset_from_predictions( art_dataset=train_dataset, y_true=val_true, y_pred=val_pred, n_classes=num_classes, train_loader_kwargs={"batch_size": 256, "shuffle": True, "drop_last": True}, )
from art_resampling import art_refresh_dataset_from_f1_scores train_loader, class_weights, per_class_f1 = art_refresh_dataset_from_f1_scores( art_dataset=train_dataset, f1=per_class_f1, train_loader_kwargs={"batch_size": 256, "shuffle": True, "drop_last": True}, )
import numpy as np from art_resampling import per_class_f1, art_weights_from_f1 val_true = np.array([0, 0, 0, 1, 1, 2, 2, 2]) val_pred = np.array([0, 0, 1, 1, 2, 0, 2, 2]) f1 = per_class_f1(val_true, val_pred, n_classes=3) weights = art_weights_from_f1(f1) print("per_class_f1:", np.round(f1, 4)) print("class_weights:", np.round(weights, 4)) print("sum:", float(weights.sum()))
pytest -q
- This library focuses on the ART resampling logic only.
- The example uses a synthetic tabular dataset for a complete end-to-end demonstration.
You choose two key ART parameters:
c: fraction of each refreshed training set drawn uniformly from all samples (the "balanced" portion).bf: refresh frequency in epochs (refresh everybfepochs).
A typical training loop looks like:
from torch.utils.data import DataLoader from art_resampling import ARTDataset, art_refresh_dataset_from_predictions c = 0.3 bf = 4 initial_class_weights = [1 / num_classes] * num_classes train_dataset = ARTDataset(X_train, y_train, enable_art=True, c=c, cls_weights=initial_class_weights, seed=seed) val_dataset = ARTDataset(X_val, y_val, enable_art=False) train_loader_kwargs = {"batch_size": 256, "shuffle": True, "drop_last": True} val_loader_kwargs = {"batch_size": 512, "shuffle": False, "drop_last": False} train_loader = DataLoader(train_dataset, **train_loader_kwargs) val_loader = DataLoader(val_dataset, **val_loader_kwargs) for epoch in range(1, epochs + 1): train_one_epoch(model, train_loader, optimizer, loss_fn) if epoch % bf == 0: val_true, val_pred = predict_labels(model, val_loader) train_loader, class_weights, per_class_f1 = art_refresh_dataset_from_predictions( art_dataset=train_dataset, y_true=val_true, y_pred=val_pred, n_classes=num_classes, train_loader_kwargs=train_loader_kwargs, )
If you find a bug, please open a GitHub Issue with:
- a minimal code snippet to reproduce
- your Python / PyTorch / numpy versions
- the exact error message or unexpected behavior
For questions about how to integrate ART into a specific training loop, include a small sketch of your dataloader and evaluation step.