Image model training

Deep dive

Decision table

Classification recipe

1. Train-time augmentation: RandomResizedCrop to 224 (area 8-100%, aspect 3/4-4/3), horizontal flip. For ViTs add RandAugment (2 ops, magnitude 9), Mixup (α 0.8), CutMix (α 1.0), random erasing (p 0.25), repeated augmentation.
2. Eval: resize the short side to 256, centre crop 224 (ratio 0.875).
3. Loss: cross-entropy with label smoothing 0.1, or BCE on soft Mixup targets ("ResNet strikes back" A1: 600 epochs, LAMB, BCE, 80.4% for ResNet-50).
4. Regularization: no weight decay on norm parameters and biases; zero-init the last BN γ of each residual block; stochastic depth 0.1 for deep ViTs.
5. Schedule: linear warmup 5 epochs, then cosine to about zero (warmup + cosine).
6. Train/test resolution: RandomResizedCrop magnifies objects, so testing at a somewhat higher resolution than training, or fine-tuning briefly at the test resolution (FixRes), gains accuracy. Recompute BN statistics at the new resolution.
7. Progressive resizing: train early epochs at low resolution (128-160) and finish at full size; EfficientNetV2 raises augmentation strength together with resolution. Most of the compute moves to cheap low-resolution steps.

Diffusion recipe

1. Data: filter by resolution, aesthetics and watermark score; dedupe; recaption with a vision-language model (DALL-E 3 trained on synthetic captions; SD3 mixed 50% original and 50% synthetic captions).
2. Pre-encode VAE latents and text embeddings once; store a flipped copy if you use horizontal flips (flipping a latent is not exactly flipping the image).
3. Resolution schedule: pretrain at 256² (fast, most steps), then 512², then about 1024² with aspect buckets. SDXL: 600k steps at 256² and 200k at 512² (batch 2048), then multi-aspect training near 1024².
4. Size and crop conditioning (SDXL): feed the original image size and crop offsets as extra conditioning so small or cropped training images can be used without teaching the model to produce small or cropped outputs.
5. High resolution: shift timesteps toward more noise (SD3: t′ = α·t/(1 + (α−1)·t) with α = √(m/n) for m vs n pixels); interpolate or extend position embeddings for DiT; QK-norm keeps attention logits bounded in mixed precision.
6. EMA: decay 0.9999 sampled at evaluation; FID on 10k-50k samples against fixed reference statistics.

Code

Building blocks on plain slices: RandomResizedCrop boxes (torchvision's algorithm), Mixup/CutMix with the clipped-box λ, soft labels with smoothing, EMA with warmup, and SDXL-style aspect buckets.

// Image-training plumbing on plain slices: RNG + Beta sampling, RandomResizedCrop boxes,
// Mixup/CutMix with soft labels, weight EMA, and aspect-ratio buckets for diffusion.
pub struct Rng(pub u64);
impl Rng {
    pub fn f(&mut self) -> f64 { // splitmix64 -> (0, 1)
        self.0 = self.0.wrapping_add(0x9E3779B97F4A7C15);
        let mut z = self.0;
        z = (z ^ (z >> 30)).wrapping_mul(0xBF58476D1CE4E5B9);
        z = (z ^ (z >> 27)).wrapping_mul(0x94D049BB133111EB);
        ((z ^ (z >> 31)) >> 11) as f64 / (1u64 << 53) as f64 + 1e-300
    }
    fn normal(&mut self) -> f64 { (-2.0 * self.f().ln()).sqrt() * (std::f64::consts::TAU * self.f()).cos() }
    fn gamma(&mut self, a: f64) -> f64 { // Marsaglia-Tsang; boost a < 1 via U^(1/a)
        if a < 1.0 { return self.gamma(a + 1.0) * self.f().powf(1.0 / a); }
        let (d, c) = (a - 1.0 / 3.0, 1.0 / (9.0 * (a - 1.0 / 3.0)).sqrt());
        loop {
            let x = self.normal();
            let v = (1.0 + c * x).powi(3);
            if v > 0.0 && self.f().ln() < 0.5 * x * x + d - d * v + d * v.ln() { return d * v; }
        }
    }
    pub fn beta(&mut self, a: f64) -> f64 { let (x, y) = (self.gamma(a), self.gamma(a)); x / (x + y) }
}

/// torchvision RandomResizedCrop box (top, left, h, w): area U(0.08, 1), log-uniform aspect 3/4..4/3.
pub fn rrc_box(rng: &mut Rng, h: usize, w: usize) -> (usize, usize, usize, usize) {
    let area = (h * w) as f64;
    for _ in 0..10 {
        let target = area * (0.08 + 0.92 * rng.f());
        let ratio = ((3f64 / 4.0).ln() + rng.f() * ((4f64 / 3.0).ln() - (3f64 / 4.0).ln())).exp();
        let (cw, ch) = ((target * ratio).sqrt().round() as usize, (target / ratio).sqrt().round() as usize);
        if cw > 0 && ch > 0 && cw <= w && ch <= h {
            return ((rng.f() * (h - ch + 1) as f64) as usize, (rng.f() * (w - cw + 1) as f64) as usize, ch, cw);
        }
    }
    let s = h.min(w); // fallback: centre crop (torchvision also clamps the ratio here)
    ((h - s) / 2, (w - s) / 2, s, s)
}

/// Mixup (cutmix = false) or CutMix of image a with b in place ([c, h, w]); returns lambda for the label mix.
pub fn mix(rng: &mut Rng, a: &mut [f32], b: &[f32], c: usize, h: usize, w: usize, alpha: f64, cutmix: bool) -> f32 {
    let lam = rng.beta(alpha);
    if !cutmix {
        for (x, y) in a.iter_mut().zip(b) { *x = lam as f32 * *x + (1.0 - lam as f32) * y; }
        return lam as f32;
    }
    let r = (1.0 - lam).sqrt(); // box side ratio so that box area = (1 - lam) before clipping
    let (cy, cx) = ((rng.f() * h as f64) as isize, (rng.f() * w as f64) as isize);
    let (bh, bw) = ((r * h as f64) as isize / 2, (r * w as f64) as isize / 2);
    let (y0, y1) = ((cy - bh).max(0) as usize, ((cy + bh) as usize).min(h));
    let (x0, x1) = ((cx - bw).max(0) as usize, ((cx + bw) as usize).min(w));
    for ch in 0..c { for y in y0..y1 { for x in x0..x1 { a[(ch * h + y) * w + x] = b[(ch * h + y) * w + x]; } } }
    1.0 - ((y1 - y0) * (x1 - x0)) as f32 / (h * w) as f32 // exact lambda after clipping
}

/// Soft target: lam * onehot(ya) + (1 - lam) * onehot(yb), plus label smoothing eps.
pub fn soft_label(n_cls: usize, ya: usize, yb: usize, lam: f32, eps: f32) -> Vec<f32> {
    let mut t = vec![eps / n_cls as f32; n_cls];
    t[ya] += (1.0 - eps) * lam;
    t[yb] += (1.0 - eps) * (1.0 - lam);
    t
}

/// EMA of weights; horizon about 1/(1 - decay) steps. Warmup (1+n)/(10+n) avoids averaging in init.
pub fn ema_update(ema: &mut [f32], w: &[f32], decay: f32, step: u64) {
    let d = decay.min((1 + step) as f32 / (10 + step) as f32);
    for (e, x) in ema.iter_mut().zip(w) { *e = d * *e + (1.0 - d) * x; }
}

/// SDXL-style buckets: sides multiple of 64, at most `px` pixels, aspect 1:2..2:1; pick by nearest log aspect.
pub fn buckets(px: usize) -> Vec<(usize, usize)> {
    (8..=32).map(|k| k * 64).map(|w| (w, (px / w) / 64 * 64)).filter(|&(w, h)| w <= 2 * h && h <= 2 * w).collect()
}
pub fn pick_bucket(b: &[(usize, usize)], w: usize, h: usize) -> (usize, usize) {
    let ar = (w as f64 / h as f64).ln();
    *b.iter().min_by(|p, q| {
        let d = |t: &(usize, usize)| ((t.0 as f64 / t.1 as f64).ln() - ar).abs();
        d(p).partial_cmp(&d(q)).unwrap()
    }).unwrap()
}

fn main() {
    let mut rng = Rng(3);
    let m: f64 = (0..20000).map(|_| rng.beta(0.8)).sum::<f64>() / 20000.0;
    assert!((m - 0.5).abs() < 0.02); // Beta(a, a) has mean 0.5
    for _ in 0..1000 { let (t, l, h, w) = rrc_box(&mut rng, 375, 500); assert!(t + h <= 375 && l + w <= 500); }
    let (mut a, b) = (vec![1.0f32; 3 * 32 * 32], vec![0.0f32; 3 * 32 * 32]);
    let lam = mix(&mut rng, &mut a, &b, 3, 32, 32, 1.0, true);
    assert!((a.iter().sum::<f32>() / a.len() as f32 - lam).abs() < 1e-6); // pixels and label agree
    let t = soft_label(10, 2, 7, lam, 0.1);
    assert!((t.iter().sum::<f32>() - 1.0).abs() < 1e-6);
    let mut e = vec![0.0f32; 4];
    for s in 0..50_000 { ema_update(&mut e, &[1.0; 4], 0.9999, s); }
    assert!(e[0] > 0.99);
    let bk = buckets(1024 * 1024);
    assert!(bk.contains(&(1024, 1024)) && pick_bucket(&bk, 1920, 1080).0 > 1024);
    println!("ok {:?}", pick_bucket(&bk, 1920, 1080));
}

Data loading and throughput

• JPEG decode is often the bottleneck: order of a few hundred images/s per CPU core at ImageNet sizes, so a fast GPU needs dozens of cores or GPU decode (NVIDIA DALI with nvJPEG). Store data as large sequential shards (WebDataset tar, FFCV) instead of millions of small files. See data loading pipelines.
• bf16 autocast, channels-last memory format for CNNs, fused optimizers, compiled graphs; see mixed precision.
• Diffusion: precomputed latents make loading trivial; the text encoder and VAE stay off the training GPUs.

Efficiency

• Cost: ResNet-50 reaches 77.2% top-1 in 40 epochs, 35.6 min on 8 A100 with FFCV (its README); batch 8192 on 256 GPUs takes 1 hour (Goyal et al.). Bound: JPEG decode and loading, else compute.
• Levers:
  • low resolution early, full size at the end — EfficientNetV2 trains 5-11x faster than compared models (its abstract) — classifiers
  • sequential shards, fast loader — keeps the GPU fed — small-file datasets
  • pre-encoded latents and text embeddings — encoders off the GPUs — diffusion
  • large batch, LR scaling, warmup — 256 GPUs (Goyal et al.) — up to about 8k, then LARS/LAMB
• Naive → standard → state of the art:
  • fp32, small JPEGs, full resolution · loader-bound
  • recipe above, one 8-GPU node · hours per run
  • fast loader, low-resolution schedule · 35.6 min for 77.2%

Pitfalls

• Mixup/CutMix labels computed from the sampled λ instead of the clipped box area: the label does not match the pixels.
• Applying Mixup/CutMix or RandAugment at evaluation, or forgetting to switch BN to eval mode.
• Evaluating the raw weights when the EMA copy is the one that is good (or vice versa early in training, when the EMA still contains the random init; use EMA warmup).
• Centre-cropping diffusion training data: generated subjects get cut at the frame edge.
• Resolution jump without timestep shift: high-resolution samples come out with incoherent global structure.
• Data loader starving the GPU: check GPU utilization before tuning the model.
• Linear LR scaling past the point where it holds (beyond about 8k for SGD on ImageNet) without LARS/LAMB.

Implementations

Sources

• Goyal et al., Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour, 2017. paper
• He et al., Bag of Tricks for Image Classification with Convolutional Neural Networks, CVPR 2019. paper
• Touvron et al., Training data-efficient image transformers & distillation through attention, ICML 2021. paper · code
• Wightman, Touvron, Jégou, ResNet strikes back: An improved training procedure in timm, 2021. paper · code
• Touvron et al., Fixing the train-test resolution discrepancy, NeurIPS 2019. paper
• Tan and Le, EfficientNetV2: Smaller Models and Faster Training, ICML 2021. paper
• Zhang et al., mixup: Beyond Empirical Risk Minimization, ICLR 2018. paper
• Yun et al., CutMix: Regularization Strategy to Train Strong Classifiers with Localizable Features, ICCV 2019. paper
• Cubuk et al., RandAugment: Practical automated data augmentation with a reduced search space, 2020. paper
• You, Gitman, Ginsburg, Large Batch Training of Convolutional Networks, 2017. paper
• Peebles and Xie, Scalable Diffusion Models with Transformers, ICCV 2023. paper · code
• Podell et al., SDXL: Improving Latent Diffusion Models for High-Resolution Image Synthesis, 2023. paper · code
• Esser et al., Scaling Rectified Flow Transformers for High-Resolution Image Synthesis, ICML 2024. paper
• Betker et al., Improving Image Generation with Better Captions, 2023. paper
• Karras et al., Analyzing and Improving the Training Dynamics of Diffusion Models, CVPR 2024. paper · code