feat: enhance ResNet examples with gradient accumulation and updated defaults

- Added `grads_accumulation` argument to support gradient accumulation for better memory management.
- Updated default values for learning rate, LR decay, batch size, drop block probability, and drop path probability.
- Refactored `LogConfig` for streamlined documentation and improved modularity.
- Improved stochastic depth initialization logic to enforce consistency within ResNet layers.

Author: crutcher

crates/bimm/src/models/resnet/layer_block.rs

@@ -243,16 +243,13 @@ impl LayerBlockStructureConfig {
         O: Into<Option<DropBlockOptions>>,
     {
         let options = options.into();
-        self.map_blocks(&mut |_, block| block.with_drop_block(options.clone()))
-    }
-
-    /// Update the drop path probability.
-    pub fn with_drop_path_prob(
-        self,
-        prob: f64,
-    ) -> Self {
-        let prob = expect_probability(prob);
-        self.map_blocks(&mut |_, block| block.with_drop_path_prob(prob))
+        self.map_blocks(&mut |idx, block| {
+            if idx == 0 {
+                block.with_drop_block(None)
+            } else {
+                block.with_drop_block(options.clone())
+            }
+        })
     }
 }
 

crates/bimm/src/models/resnet/resnet_model.rs

@@ -214,13 +214,13 @@ impl ResNetStructureConfig {
     ) -> Self {
         let drop_path_rate = expect_probability(drop_path_rate);
 
-        let net_num_blocks = self.layers.iter().map(|b| b.len()).sum::<usize>();
+        let net_num_blocks = self.layers.iter().map(|b| b.len()).sum::<usize>() - self.layers.len();
         let mut net_block_idx = 0;
-        let mut update_drop_path = |_idx: usize, block: ResidualBlockStructureConfig| {
+        let mut update_drop_path = |idx: usize, block: ResidualBlockStructureConfig| {
             // stochastic depth linear decay rule
             let block_dpr = drop_path_rate * (net_block_idx as f64) / ((net_num_blocks - 1) as f64);
             net_block_idx += 1;
-            if block_dpr > 0.0 {
+            if idx != 0 && block_dpr > 0.0 {
                 block.with_drop_path_prob(block_dpr)
             } else {
                 block

examples/resnet-finetune/src/main.rs

@@ -75,11 +75,15 @@ pub struct Args {
     artifact_dir: String,
 
     /// Batch size for processing
-    #[arg(short, long, default_value_t = 32)]
+    #[arg(short, long, default_value_t = 24)]
     batch_size: usize,
 
+    /// Grads accumulation size for processing
+    #[arg(short, long, default_value_t = 8)]
+    grads_accumulation: usize,
+
     /// Number of workers for data loading.
-    #[arg(long, default_value = "2")]
+    #[arg(long, default_value = "4")]
     num_workers: usize,
 
     /// Number of epochs to train the model.
@@ -103,55 +107,62 @@ pub struct Args {
     drop_path_prob: f64,
 
     /// Learning rate
-    #[arg(long, default_value = "1e-5")]
+    #[arg(long, default_value_t = 5e-5)]
     pub learning_rate: f64,
 
     /// Early stopping patience
-    #[arg(long, default_value = "10")]
+    #[arg(long, default_value_t = 10)]
     patience: usize,
 
     /// Optimizer Weight decay.
     #[arg(long, default_value_t = 5e-4)]
     pub weight_decay: f32,
 }
 
-/// Log config.
-///
-/// Only exists for logging.
-#[derive(Config, Debug)]
 #[allow(clippy::too_many_arguments)]
-pub struct LogConfig {
-    seed: u64,
-    train_percentage: u8,
-    batch_size: usize,
-    num_epochs: usize,
-    resnet_prefab: String,
-    resnet_pretrained: String,
-    drop_block_prob: f64,
-    drop_path_prob: f64,
-    learning_rate: f64,
-    patience: usize,
-    weight_decay: f32,
-    resnet: ResNetContractConfig,
+mod local {
+    use bimm::models::resnet::ResNetContractConfig;
+    use burn::config::Config;
+
+    /// Log config.
+    ///
+    /// Only exists for logging.
+    #[derive(Config, Debug)]
+    pub struct LogConfig {
+        pub seed: u64,
+        pub train_percentage: u8,
+        pub batch_size: usize,
+        pub num_epochs: usize,
+        pub resnet_prefab: String,
+        pub resnet_pretrained: String,
+        pub drop_block_prob: f64,
+        pub drop_path_prob: f64,
+        pub learning_rate: f64,
+        pub patience: usize,
+        pub weight_decay: f32,
+        pub resnet: ResNetContractConfig,
+    }
 }
+use local::*;
 
-fn main() {
+fn main() -> anyhow::Result<()> {
     let args = Args::parse();
 
     let _source_tree = download();
 
     let device = Default::default();
-    train::<Autodiff<Cuda>>(&args, &device);
+    train::<Autodiff<Cuda>>(&args, &device)
 }
 
+#[must_use]
 pub fn train<B: AutodiffBackend>(
     args: &Args,
     device: &B::Device,
 ) -> anyhow::Result<()> {
     // Remove existing artifacts before to get an accurate learner summary
     let artifact_dir: &str = args.artifact_dir.as_ref();
-    std::fs::remove_dir_all(artifact_dir)?;
-    std::fs::create_dir_all(artifact_dir)?;
+    std::fs::remove_dir_all(artifact_dir);
+    std::fs::create_dir_all(artifact_dir).expect("Failed to create artifacts directory");
 
     B::seed(args.seed);
 
@@ -232,6 +243,7 @@ pub fn train<B: AutodiffBackend>(
             },
         ))
         .devices(vec![device.clone()])
+        .grads_accumulation(args.grads_accumulation)
         .num_epochs(args.num_epochs)
         .summary()
         .build(model, optimizer, args.learning_rate);

examples/resnet_tiny/src/main.rs

@@ -66,6 +66,10 @@ pub struct Args {
     #[arg(short, long, default_value_t = 512)]
     batch_size: usize,
 
+    /// Grads accumulation size for processing
+    #[arg(short, long, default_value_t = 8)]
+    grads_accumulation: usize,
+
     /// Number of workers for data loading.
     #[arg(long, default_value = "2")]
     num_workers: Option<usize>,
@@ -79,11 +83,11 @@ pub struct Args {
     drop_block_rate: f64,
 
     /// Learning rate for the optimizer.
-    #[arg(long, default_value = "1.0e-6")]
+    #[arg(long, default_value = "1.0e-4")]
     learning_rate: f64,
 
     /// Learning rate decay gamma.
-    #[arg(long, default_value = "0.999975")]
+    #[arg(long, default_value = "0.999997")]
     lr_gamma: f64,
 
     /// Directory to save the artifacts.
@@ -107,11 +111,11 @@ pub struct Args {
     resnet_pretrained: Option<String>,
 
     /// Drop Block Prob
-    #[arg(long, default_value = "0.25")]
+    #[arg(long, default_value = "0.20")]
     drop_block_prob: f64,
 
     /// Drop Path Prob
-    #[arg(long, default_value = "0.15")]
+    #[arg(long, default_value = "0.0")]
     drop_path_prob: f64,
 
     /// Early stopping patience
@@ -315,6 +319,7 @@ pub fn backend_main<B: AutodiffBackend>(
             },
         ))
         .devices(devices.clone())
+        .grads_accumulation(args.grads_accumulation)
         .num_epochs(args.num_epochs)
         .summary()
         .build(model, optim_config.init(), lr_scheduler);