Data augmentation

Deep dive

How it works

In: example (x, y, optional boxes/masks/keypoints), seed, epoch, index · Out: transformed example with consistent targets · State: none (RNG seeded from (seed, epoch, index)) · Cost: O(input size) per transform, usually CPU-bound on decode.

1. Pick transforms that keep the label. Per domain: images (crop, flip, rotation, colour), audio and signals (time shift, noise, gain, time/frequency masking), text (synonym swap, back-translation), medical (small rotation, intensity), tabular (interpolate between minority neighbours).
2. Seed per sample. rng = f(seed, epoch, index), so worker count and order do not change the data.
3. Sample parameters (offset, flip, noise level), uniformly or from a policy: AutoAugment learns sub-policies on validation accuracy; RandAugment picks N random ops at one shared magnitude M, tuned directly on the target task.
4. Apply geometric ops to input and targets together. Shift or flip moves a box with it; clip it to the window and drop it if empty. Photometric ops (noise, brightness) touch only x.
5. Optionally mix. mixup: x = λxa + (1-λ)xb, same for labels, λ ~ Beta(α,α). CutMix: paste a patch from b, mix labels by area.
6. Train on the result; evaluate on clean data. Test-time augmentation (TTA) averages predictions over a few variants.

Code

Seeded 1-D sensor pipeline: shift (box moves), optional time reversal (box mirrors), noise, mixup (α = 1, box dropped), with consistency and reproducibility asserts.

/// splitmix64 stream; one per sample, seeded from (seed, epoch, index) so any worker order gives the same data.
struct Rng(u64);
impl Rng {
    fn new(seed: u64, epoch: u64, idx: u64) -> Self {
        Rng(seed ^ epoch.wrapping_mul(0x9E37_79B9_7F4A_7C15) ^ idx.wrapping_mul(0xD1B5_4A32_D192_ED03))
    }
    fn next(&mut self) -> u64 {
        self.0 = self.0.wrapping_add(0x9E37_79B9_7F4A_7C15);
        let mut z = self.0;
        z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
        z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
        z ^ (z >> 31)
    }
    fn unit(&mut self) -> f32 { (self.next() >> 40) as f32 / (1u64 << 24) as f32 } // [0,1)
    fn below(&mut self, n: usize) -> usize { (self.next() % n as u64) as usize }
    fn normal(&mut self) -> f32 { // Box-Muller
        let (u, v) = (self.unit().max(1e-7), self.unit());
        (-2.0 * u.ln()).sqrt() * (6.2831855 * v).cos()
    }
}
#[derive(Clone)]
struct Sample { x: Vec<f32>, label: [f32; 2], span: Option<(usize, usize)> } // soft label + event box [s,e)
/// Geometric ops move the box with the signal; photometric ops leave it alone.
fn augment(s: &Sample, r: &mut Rng, pad: usize, flip_ok: bool, sigma: f32) -> Sample {
    let n = s.x.len();
    let off = r.below(2 * pad + 1) as isize - pad as isize; // random shift, zero fill
    let mut x = vec![0.0; n];
    for i in 0..n {
        let j = i as isize - off;
        if j >= 0 && (j as usize) < n { x[i] = s.x[j as usize]; }
    }
    let mut span = s.span.and_then(|(a, b)| {
        let (a, b) = ((a as isize + off).max(0) as usize, (b as isize + off).min(n as isize).max(0) as usize);
        (b > a).then_some((a, b)) // box shifted out of the window: no box left
    });
    if flip_ok && r.next() & 1 == 1 { // time reversal; only if the label is reversal-invariant
        x.reverse();
        span = span.map(|(a, b)| (n - b, n - a));
    }
    for v in x.iter_mut() { *v += sigma * r.normal(); }
    Sample { x, label: s.label, span }
}
/// mixup with alpha=1 (lambda ~ U(0,1)): convex combination of inputs and labels; boxes are dropped (ill-defined).
fn mixup(a: &Sample, b: &Sample, r: &mut Rng) -> Sample {
    let l = r.unit();
    Sample {
        x: a.x.iter().zip(&b.x).map(|(p, q)| l * p + (1.0 - l) * q).collect(),
        label: [l * a.label[0] + (1.0 - l) * b.label[0], l * a.label[1] + (1.0 - l) * b.label[1]],
        span: None,
    }
}
fn batch(data: &[Sample], seed: u64, epoch: u64, ids: &[usize]) -> Vec<Sample> {
    ids.iter().map(|&i| {
        let mut r = Rng::new(seed, epoch, i as u64);
        let a = augment(&data[i], &mut r, 16, true, 0.05);
        let j = r.below(data.len());
        let b = augment(&data[j], &mut r, 16, true, 0.05);
        mixup(&a, &b, &mut r)
    }).collect()
}
fn pulse(n: usize, s: usize, e: usize, c: usize) -> Sample {
    let mut x = vec![0.0; n];
    x[s..e].iter_mut().for_each(|v| *v = 1.0);
    Sample { x, label: if c == 0 { [1.0, 0.0] } else { [0.0, 1.0] }, span: Some((s, e)) }
}
fn main() {
    let n = 256;
    let data: Vec<Sample> = (0..64).map(|i| pulse(n, 40 + i, 90 + i, i % 2)).collect();
    // 1. box and signal move together (no mixup): mean inside box high, outside low
    for k in 0..200 {
        let a = augment(&data[k % 64], &mut Rng::new(7, 0, k as u64), 16, true, 0.05);
        if let Some((s, e)) = a.span {
            let inside: f32 = a.x[s..e].iter().sum::<f32>() / (e - s) as f32;
            let outside: f32 = (a.x[..s].iter().chain(&a.x[e..]).sum::<f32>()) / (n - (e - s)) as f32;
            assert!(inside > 0.8 && outside.abs() < 0.1, "box off the pulse");
        }
    }
    // 2. mixed labels sum to 1; same seed => identical batch; other epoch differs
    let ids: Vec<usize> = (0..32).collect();
    let (b1, b2, b3) = (batch(&data, 42, 0, &ids), batch(&data, 42, 0, &ids), batch(&data, 42, 1, &ids));
    assert!(b1.iter().all(|s| (s.label[0] + s.label[1] - 1.0).abs() < 1e-5));
    assert!(b1.iter().zip(&b2).all(|(p, q)| p.x == q.x && p.label == q.label));
    assert!(b1.iter().zip(&b3).any(|(p, q)| p.x != q.x));
    // throughput: 1-D windows of 1024 floats, one thread
    let big: Vec<Sample> = (0..256).map(|i| pulse(1024, 100 + i, 300 + i, i % 2)).collect();
    let ids: Vec<usize> = (0..256).collect();
    let t = std::time::Instant::now();
    let mut chk = 0.0f32;
    for ep in 0..200 { chk += batch(&big, 1, ep, &ids)[0].x[0]; }
    let dt = t.elapsed().as_secs_f64();
    println!("ok; {:.0} samples/s ({:.0} MB/s in), chk {chk}", 256.0 * 200.0 / dt, 256.0 * 200.0 * 4096.0 / dt / 1e6);
}

Measured on an Apple M4 Pro, release build, one thread: 40k samples/s (order of 165 MB/s of input) for two shift/flip/noise passes plus mixup on 1024-float windows (3 runs, 39.5k-42k).

Complexity and constants

• O(n) per sample and per transform; mixup and CutMix cost one extra read of a second sample.
• Memory-bound for simple ops, compute-bound for rotations, warps and spectrograms.
• The sampler (here ~25 µs per sample, mostly Box-Muller noise) is cheap next to decoding a JPEG or resampling audio.

Efficiency

• Cost: on this M4 Pro one thread does about 40k samples/s of 1024 floats (measured). For images the bound is CPU decode plus resize/warp (order of ms per image), not the RNG; it competes with the training step for cores.
• Levers:
  • Do augmentation on the GPU after one decode (Kornia, DALI, torchvision v2 on tensors) — frees CPU cores — when the loader is the bottleneck.
  • Cache the decoded, resized dataset and augment cheap ops online — skips repeated decode — when the dataset fits in RAM.
  • Offline augmentation (write variants to disk) — fixed cost once, but variants repeat every epoch and storage multiplies — only for expensive ops (back-translation, SMOTE).
  • Fuse ops (crop+flip+normalize in one pass) — fewer memory round trips — for large inputs.
  • Seed per sample — reproducible runs and exact resume (see Data loading pipelines) — always.
• Naive → standard → state of the art: precompute every variant on disk · storage × K, frozen variants; online per-sample CPU transforms in worker processes · cost scales with cores; decode once then batched GPU augmentation · CPU freed, augmentation cost small next to the model step.

Implementation notes

• Transform on the smallest representation: crop before resize, mask on the spectrogram (STFT and spectrogram), not on the waveform.
• Sample all parameters first, then apply, so the same parameters hit image, mask and boxes.
• Mixed labels need a soft-label loss (cross-entropy on probabilities); mixup makes the model underconfident, so recheck Calibration.
• Boxes after a shift or crop: clip, then drop boxes that lose most of their area (Object detection).

Pitfalls

• Label-changing transforms. Flipping digits (6 vs 9, 2 vs 5), rotating text or maps, mirroring left/right in medical images (heart side, lateral markers), time-reversing a sensor trace with directional events.
• Leakage. Augment copies after splitting. Augmenting first, or oversampling (SMOTE) before the split, puts near-copies of test items in training and inflates Evaluation metrics. Split by patient, speaker or session first.
• Augmenting evaluation data except deliberately for TTA.
• Train/test mismatch. Heavy noise or masking in training with a clean test is usually fine; a train pipeline that differs from inference preprocessing (resize, normalization) is not.
• Oversampling a rare class by duplication or SMOTE raises the base rate the model sees: recalibrate scores.
• Over-strong augmentation hurts small models; RandAugment's M is the dial.
• Non-reproducible RNG shared across workers gives duplicate or order-dependent batches.

Variants and alternatives

In production

• Vision: torchvision and albumentations pipelines in classification and detection training; RandAugment used in ImageNet recipes.
• Speech: SpecAugment in end-to-end ASR (LibriSpeech 6.8% WER on test-other without a language model, per the paper).
• Medical imaging, sensors and fraud: geometric/intensity transforms, jitter and time-shift, class oversampling, each restricted to what the clinical or physical label allows.

Implementations

Sources

• Zhang, Cisse, Dauphin, Lopez-Paz, mixup: Beyond Empirical Risk Minimization, ICLR 2018. paper
• Yun et al., CutMix: Regularization Strategy to Train Strong Classifiers with Localizable Features, ICCV 2019. paper · code
• Cubuk et al., AutoAugment: Learning Augmentation Policies from Data, 2018. paper
• Cubuk, Zoph, Shlens, Le, RandAugment: Practical automated data augmentation with a reduced search space, CVPR Workshops 2020. paper
• Park et al., SpecAugment: A Simple Data Augmentation Method for Automatic Speech Recognition, Interspeech 2019. paper
• Chawla, Bowyer, Hall, Kegelmeyer, SMOTE: Synthetic Minority Over-sampling Technique, JAIR 16:321-357, 2002. paper
• Shorten, Khoshgoftaar, A survey on Image Data Augmentation for Deep Learning, Journal of Big Data 6:60, 2019. paper