$ show google-research/augmix
This repository contains the reference implementation of AugMix, a data augmentation technique for improving robustness and uncertainty calibration in image classification models. It provides training scripts for CIFAR-10/100 and ImageNet using PyTorch, along with evaluation on their corrupted versions (CIFAR-C, ImageNet-C). The core augmentation logic is implemented in NumPy and PIL, and the training code applies a Jensen-Shannon Divergence consistency loss between augmented and clean images. The codebase is research-oriented, with scripts that are well-structured but heavily tied to the specific experimental setup of the paper. Dependencies are pinned but severely outdated (torch==1.2.0), making reproduction on modern systems challenging without adaptation.
The LLM identified 3 files worth highlighting for reuse:
Provides a collection of PIL-based image augmentation operators (autocontrast, equalize, posterize, rotate, solarize, shear, translate, color, contrast, brightness, sharpness) with severity-driven parameter scaling. Each operator is a standalone function that takes a PIL image and a severity level, making them reusable across any image augmentation pipeline. The parameter scaling pattern (int_parameter, float_parameter, sample_level) is clean and easily adaptable to new operations.
Implements the core AugMix algorithm in numpy: combines multiple augmentation chains via Dirichlet-distributed weights and a convex mixture with the original image. The function 'augment_and_mix' is well-parameterized (severity, width, depth, alpha) and operates on numpy arrays, making it model-agnostic and easily extractable into other projects (e.g., TensorFlow, JAX) with minimal adaptation.
Contains the 'compute_mce' function that calculates the mean Corruption Error (mCE) normalized by AlexNet baseline errors. This evaluation metric is standard in robustness research and is implemented in a clean, reusable manner. It takes a dictionary of corruption accuracies and returns a single mCE score. Could be extracted as a standalone utility for benchmarking model robustness on ImageNet-C.