[WIP] Add timm inspired burn_vision::layers::drop::{drop_path, DropPath}

crates/burn-vision/Cargo.toml

@@ -36,6 +36,7 @@ test-wgpu = ["cubecl-backend", "export-tests"]
 
 [dependencies]
 bon = { workspace = true }
+burn-core = { path = "../burn-core", version = "0.18.0"}
 burn-candle = { path = "../burn-candle", version = "0.18.0", optional = true }
 burn-cubecl = { path = "../burn-cubecl", version = "0.18.0", optional = true }
 burn-fusion = { path = "../burn-fusion", version = "0.18.0", optional = true }

crates/burn-vision/src/layers/drop.rs

@@ -0,0 +1,182 @@
+/// Burn implementation of the DropPath (Stochastic Depth) regularization layer.
+///
+/// Papers:
+/// DropBlock: A regularization method for convolutional networks (https://arxiv.org/abs/1810.12890)
+///
+/// Deep Networks with Stochastic Depth (https://arxiv.org/abs/1603.09382)
+///
+/// Inspired by the python implementation from the timm library:
+/// https://github.com/huggingface/pytorch-image-models/blob/main/timm/layers/drop.py
+use burn_core as burn;
+use burn_core::config::Config;
+use burn_core::module::Module;
+use burn_tensor::backend::Backend;
+use burn_tensor::{Distribution, Tensor};
+
+fn check_probability(prob: f64) -> f64 {
+    if !(0.0..=1.0).contains(&prob) {
+        panic!("Probability should be between 0 and 1, but got {}", prob);
+    }
+    prob
+}
+
+/// DropPath (stochastic depth) regularization.
+///
+/// ## Arguments
+///
+/// * `x`: Input tensor.
+/// * `drop_prob`: Probability of dropping a path.
+/// * `training`: Whether the model is in training mode.
+/// * `scale_by_keep`: Whether to scale the output by `1 / (1 - drop_prob)`
+///
+/// ## Returns
+///
+/// * Output tensor with the same shape as the input tensor.
+#[must_use]
+pub fn drop_path<B: Backend, const D: usize>(
+    x: Tensor<B, D>,
+    drop_prob: f64,
+    training: bool,
+    scale_by_keep: bool,
+) -> Tensor<B, D> {
+    _drop_path_sample(
+        x,
+        drop_prob,
+        training,
+        scale_by_keep,
+        |shape, keep_prob, device| {
+            Tensor::<B, D>::random(shape, Distribution::Bernoulli(keep_prob), device)
+        },
+    )
+}
+
+/// Internal implementation of DropPath.
+///
+/// Deferred to a separate function to allow for testing sampling.
+///
+/// ## Arguments
+///
+/// * `x`: Input tensor.
+/// * `drop_prob`: Probability of dropping a path.
+/// * `training`: Whether the model is in training mode.
+/// * `scale_by_keep`: Whether to scale the output by `1 / (1 - drop_prob)`
+/// * `sample`: Sampling function to generate the random tensor.
+///
+/// ## Returns
+///
+/// * Output tensor with the same shape as the input tensor.
+#[inline(always)]
+#[must_use]
+fn _drop_path_sample<B: Backend, const D: usize>(
+    x: Tensor<B, D>,
+    drop_prob: f64,
+    training: bool,
+    scale_by_keep: bool,
+    sample: fn([usize; D], f64, &B::Device) -> Tensor<B, D>,
+) -> Tensor<B, D> {
+    check_probability(drop_prob);
+
+    if !training || drop_prob == 0.0 {
+        return x;
+    }
+
+    let keep_prob = 1.0 - drop_prob;
+
+    let mut shape = [1; D];
+    shape[0] = x.dims()[0];
+
+    let random_tensor = sample(shape, keep_prob, &x.device());
+
+    let random_tensor = if keep_prob > 0.0 && scale_by_keep {
+        random_tensor.div_scalar(keep_prob)
+    } else {
+        random_tensor
+    };
+
+    x * random_tensor
+}
+
+/// Common introspection interface for DropPath modules.
+pub trait DropPathMeta {
+    /// Probability of dropping a path.
+    fn drop_prob(&self) -> f64;
+
+    /// Probability of keeping a path; the complement of `drop_prob`.
+    fn keep_prob(&self) -> f64 {
+        1.0 - self.drop_prob()
+    }
+
+    /// Whether to scale the output by `1 / (1 - drop_prob)`.
+    fn scale_by_keep(&self) -> bool;
+}
+
+/// Configuration for the DropPath module.
+#[derive(Config, Debug)]
+pub struct DropPathConfig {
+    /// Probability of dropping a path.
+    #[config(default = 0.0)]
+    pub drop_prob: f64,
+
+    /// Whether to scale the output by `1 / (1 - drop_prob)`.
+    #[config(default = true)]
+    pub scale_by_keep: bool,
+}
+
+impl DropPathMeta for DropPathConfig {
+    fn drop_prob(&self) -> f64 {
+        self.drop_prob
+    }
+
+    fn scale_by_keep(&self) -> bool {
+        self.scale_by_keep
+    }
+}
+
+impl DropPathConfig {
+    /// Initializes a new DropPath module.
+    #[must_use]
+    pub fn init(&self) -> DropPath {
+        DropPath {
+            drop_prob: check_probability(self.drop_prob),
+            scale_by_keep: self.scale_by_keep,
+        }
+    }
+}
+
+/// The DropPath module.
+///
+/// Burn Module that implements the DropPath (Stochastic Depth) regularization.
+#[derive(Module, Clone, Debug)]
+pub struct DropPath {
+    /// Probability of dropping a path.
+    pub drop_prob: f64,
+
+    /// Whether to scale the output by `1 / (1 - drop_prob)`.
+    pub scale_by_keep: bool,
+}
+
+impl DropPathMeta for DropPath {
+    fn drop_prob(&self) -> f64 {
+        self.drop_prob
+    }
+
+    fn scale_by_keep(&self) -> bool {
+        self.scale_by_keep
+    }
+}
+
+impl DropPath {
+    /// Applies the forward pass on the input tensor.
+    ///
+    /// See [DropPath](DropPath) for more information.
+    ///
+    /// # Shapes
+    ///
+    /// - input: `[..., any]`
+    /// - output: `[..., any]`
+    #[must_use]
+    pub fn forward<B: Backend, const D: usize>(&self, input: Tensor<B, D>) -> Tensor<B, D> {
+        let training = B::ad_enabled();
+        drop_path(input, self.drop_prob, training, self.scale_by_keep)
+    }
+}

crates/burn-vision/src/layers/mod.rs

@@ -0,0 +1,2 @@
+/// Drop image modules.
+pub mod drop;

crates/burn-vision/src/lib.rs

@@ -14,6 +14,10 @@ extern crate alloc;
 
 /// Backend implementations for JIT and CPU
 pub mod backends;
+
+/// Image layers for burn-vision
+pub mod layers;
+
 mod base;
 mod ops;
 mod tensor;

crates/burn-vision/src/tests/layers/drop.rs

@@ -0,0 +1,133 @@
+#[burn_tensor_testgen::testgen(layers_drop)]
+mod tests {
+    use super::*;
+    use burn_vision::layers::drop::*;
+
+    #[test]
+    fn test_drop_path() {
+        let device = Default::default();
+        let drop_prob = 0.5;
+        let scale_by_keep = true;
+
+        let config = DropPathConfig {
+            drop_prob,
+            scale_by_keep,
+        };
+
+        let module = config.init();
+
+        let input = Tensor::<TestBackend, 4>::random(
+            [2, 3, 4, 5],
+            Distribution::Uniform(0.0, 1.0),
+            &device,
+        );
+        let output = module.forward(input.clone());
+
+        assert_eq!(input.dims(), output.dims());
+    }
+
+    #[test]
+    fn test_drop_path_sample() {
+        let device = Default::default();
+
+        let n = 3;
+        let shape = [n, 2, 4];
+
+        let x = Tensor::<TestBackend, 3>::random(shape, Distribution::Uniform(0.0, 1.0), &device);
+
+        /// No-op case: not training and drop_prob = 0.0
+        let training = false;
+        let drop_prob = 0.0;
+        let scale_by_keep = false;
+        let res = _drop_path_sample(
+            x.clone(),
+            drop_prob,
+            training,
+            scale_by_keep,
+            |shape, keep_prob, device| {
+                Tensor::<TestBackend, 3>::from_data([[[1.0]], [[0.0]], [[1.0]]], device)
+            },
+        );
+        res.to_data().assert_eq(&x.clone().to_data(), true);
+
+        /// No-op case: training, but drop_prob = 0.0
+        let training = true;
+        let drop_prob = 0.0;
+        let scale_by_keep = false;
+        let res = _drop_path_sample(
+            x.clone(),
+            drop_prob,
+            training,
+            scale_by_keep,
+            |shape, keep_prob, device| {
+                Tensor::<TestBackend, 3>::from_data([[[1.0]], [[0.0]], [[1.0]]], device)
+            },
+        );
+        res.to_data().assert_eq(&x.clone().to_data(), true);
+
+        /// Training, but no scaling
+        let training = true;
+        let drop_prob = 0.5;
+        let scale_by_keep = false;
+        let res = _drop_path_sample(
+            x.clone(),
+            drop_prob,
+            training,
+            scale_by_keep,
+            |shape, keep_prob, device| {
+                Tensor::<TestBackend, 3>::from_data([[[1.0]], [[0.0]], [[1.0]]], device)
+            },
+        );
+        res.to_data().assert_eq(
+            &(x.clone()
+                * Tensor::<TestBackend, 3>::from_data([[[1.0]], [[0.0]], [[1.0]]], &device))
+            .to_data(),
+            true,
+        );
+
+        /// Training, with scaling
+        let training = true;
+        let drop_prob = 0.5;
+        let keep_prob = 1.0 - drop_prob;
+        let scale_by_keep = true;
+        let res = _drop_path_sample(
+            x.clone(),
+            drop_prob,
+            training,
+            scale_by_keep,
+            |shape, keep_prob, device| {
+                Tensor::<TestBackend, 3>::from_data([[[1.0]], [[0.0]], [[1.0]]], device)
+            },
+        );
+        res.to_data().assert_eq(
+            &(x.clone()
+                * Tensor::<TestBackend, 3>::from_data([[[1.0]], [[0.0]], [[1.0]]], &device))
+            .div_scalar(keep_prob)
+            .to_data(),
+            true,
+        );
+    }
+
+    #[test]
+    fn test_droppath_module() {
+        let drop_prob = 0.2;
+        let config = DropPathConfig::new().with_drop_prob(drop_prob);
+
+        assert_eq!(config.drop_prob(), 0.2);
+        assert_eq!(config.keep_prob(), 1.0 - drop_prob);
+        assert!(config.scale_by_keep());
+
+        let module = config.init();
+        assert_eq!(module.drop_prob(), 0.2);
+        assert_eq!(module.keep_prob(), 1.0 - drop_prob);
+        assert!(module.scale_by_keep());
+
+        let device = Default::default();
+        let shape = [2, 3, 4];
+        let x = Tensor::<TestBackend, 3>::random(shape, Distribution::Uniform(0.0, 1.0), &device);
+
+        // TODO(crutcher): work out how to enable/disable training mode in tests.
+        let output = module.forward(x.clone());
+        assert_eq!(x.dims(), output.dims());
+    }
+}

crates/burn-vision/src/tests/layers/mod.rs

@@ -0,0 +1 @@
+pub(crate) mod drop;

crates/burn-vision/src/tests/mod.rs

@@ -4,6 +4,7 @@ use burn_tensor::{Shape, Tensor, TensorData, backend::Backend};
 use image::{DynamicImage, ImageBuffer, Luma, Rgb};
 
 mod connected_components;
+mod layers;
 mod morphology;
 
 #[macro_export]
@@ -22,6 +23,8 @@ macro_rules! testgen_all {
 
             burn_vision::testgen_connected_components!();
             burn_vision::testgen_morphology!();
+
+            burn_vision::testgen_layers_drop!();
         }
     };
 }