burn_telemetry 0.1.0

//! MNIST data utilities adapted from the official Burn example.
//!
//! Portions of this module are based on
//! https://github.com/tracel-ai/burn/tree/v0.19.0/examples/mnist
//! which is available under the terms of the MIT OR Apache-2.0 licenses.

use burn::{
    backend::NdArray,
    data::{
        dataloader::batcher::Batcher,
        dataset::{transform::Mapper, vision::MnistItem, Dataset},
    },
    prelude::*,
    vision::Transform2D,
};
use rand::Rng;
use std::{f32::consts::FRAC_PI_4, sync::Arc};

/// Batch type used by the MNIST loaders.
#[derive(Clone, Debug)]
pub struct MnistBatch<B: Backend> {
    pub images: Tensor<B, 3>,
    pub targets: Tensor<B, 1, Int>,
}

/// Default batcher used by the shared MNIST training harness.
#[derive(Clone, Debug, Default)]
pub struct MnistBatcher;

impl<B: Backend> Batcher<B, MnistItemPrepared, MnistBatch<B>> for MnistBatcher {
    fn batch(&self, items: Vec<MnistItemPrepared>, device: &B::Device) -> MnistBatch<B> {
        let images = items.iter().map(|item| item.image.clone()).collect();

        let targets = items
            .iter()
            .map(|item| {
                Tensor::<NdArray, 1, Int>::from_data(
                    TensorData::from([(item.label as i64).elem::<B::IntElem>()]),
                    &Default::default(),
                )
            })
            .collect();

        let images = Tensor::cat(images, 0);
        let images = Tensor::from_data(images.into_data(), device);

        let targets = Tensor::cat(targets, 0);
        let targets = Tensor::from_data(targets.into_data(), device);

        MnistBatch { images, targets }
    }
}

/// Lightweight dataset fixture for sampling batches outside the learner loop.
#[derive(Clone)]
pub struct SampleProvider<B: Backend> {
    dataset: Arc<dyn Dataset<MnistItemPrepared> + Send + Sync>,
    batcher: MnistBatcher,
    len: usize,
    device: B::Device,
}

impl<B: Backend> SampleProvider<B> {
    pub fn new(
        dataset: Arc<dyn Dataset<MnistItemPrepared> + Send + Sync>,
        device: B::Device,
    ) -> Self {
        let len = dataset.len();
        Self {
            dataset,
            batcher: MnistBatcher,
            len,
            device,
        }
    }

    /// Sample a contiguous slice of the prepared dataset starting at the provided offset.
    pub fn sample(&self, offset: usize, batch_size: usize) -> Option<MnistBatch<B>> {
        if self.len == 0 || batch_size == 0 {
            return None;
        }

        let mut items = Vec::with_capacity(batch_size);
        for idx in offset..offset + batch_size {
            let sample_idx = idx % self.len;
            if let Some(item) = self.dataset.get(sample_idx) {
                items.push(item);
            }
        }

        if items.is_empty() {
            None
        } else {
            Some(self.batcher.batch(items, &self.device))
        }
    }

    /// Deterministically sample a batch for an epoch so hooks can log comparable media.
    #[allow(dead_code)]
    pub fn sample_epoch(&self, epoch: usize, batch_size: usize) -> Option<MnistBatch<B>> {
        if self.len == 0 {
            return None;
        }
        let start = (epoch * batch_size) % self.len;
        self.sample(start, batch_size)
    }
}

/// Augmentation helpers borrowed from Burn.
#[derive(Clone, Debug, Copy)]
pub enum Transform {
    Translate,
    Shear,
    Scale,
    Rotation,
}

impl core::fmt::Display for Transform {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Transform::Translate => f.write_str("Tr"),
            Transform::Shear => f.write_str("Sr"),
            Transform::Scale => f.write_str("Sc"),
            Transform::Rotation => f.write_str("Rot"),
        }
    }
}

#[derive(Default, Clone)]
pub struct MnistMapper {
    transforms: Vec<Transform>,
}

impl MnistMapper {
    #[allow(dead_code)]
    pub fn transform(mut self, transforms: &[Transform]) -> Self {
        for t in transforms {
            self.transforms.push(*t);
        }
        self
    }
    pub fn translate(mut self) -> Self {
        self.transforms.push(Transform::Translate);
        self
    }
    pub fn shear(mut self) -> Self {
        self.transforms.push(Transform::Shear);
        self
    }
    pub fn scale(mut self) -> Self {
        self.transforms.push(Transform::Scale);
        self
    }
    pub fn rotation(mut self) -> Self {
        self.transforms.push(Transform::Rotation);
        self
    }
}

impl Mapper<MnistItem, MnistItemPrepared> for MnistMapper {
    fn map(&self, item: &MnistItem) -> MnistItemPrepared {
        prepare_image(&self.transforms, item.clone())
    }
}

#[derive(Clone, Debug)]
pub struct MnistItemPrepared {
    pub image: Tensor<NdArray, 3>,
    pub label: u8,
}

fn prepare_image(transforms: &[Transform], item: MnistItem) -> MnistItemPrepared {
    let data = TensorData::from(item.image);
    let tensor = Tensor::<NdArray, 2>::from_data(data.convert::<f32>(), &Default::default());
    let tensor = tensor.reshape([1, 28, 28]);

    // normalize: make between [0,1] and make the mean =  0 and std = 1
    // values mean=0.1307,std=0.3081 were copied from PyTorch MNIST example
    let tensor = ((tensor / 255) - 0.1307) / 0.3081;
    let tensor = if !transforms.is_empty() {
        mangle_image_batch(transforms, tensor)
    } else {
        tensor
    };

    MnistItemPrepared {
        image: tensor,
        label: item.label,
    }
}

/// Apply small random transformations to augment the dataset.
fn mangle_image_batch<B: Backend>(transforms: &[Transform], images: Tensor<B, 3>) -> Tensor<B, 3> {
    let mut rng = rand::rng();

    let transforms = transforms.iter().map(|transform| match transform {
        Transform::Translate => {
            Transform2D::translation(rng.random_range(-0.2..0.2), rng.random_range(-0.2..0.2))
        }
        Transform::Shear => Transform2D::shear(
            rng.random_range(-0.6..0.6),
            rng.random_range(-0.6..0.6),
            0.0,
            0.0,
        ),
        Transform::Scale => Transform2D::scale(
            rng.random_range(0.6..1.5),
            rng.random_range(0.6..1.5),
            0.0,
            0.0,
        ),
        Transform::Rotation => {
            Transform2D::rotation(rng.random_range(-FRAC_PI_4..FRAC_PI_4), 0.0, 0.0)
        }
    });

    Transform2D::composed(transforms)
        .transform(images.unsqueeze_dim::<4>(1))
        .squeeze_dims::<3>(&[1])
}