Fix bottleneck support. (#99)

* ResNet50 WIP

* feat: add label smoothing and refactor ResNet-finetune model

- Introduced optional label smoothing for training via `--smoothing` argument.
- Refactored `ResNet` into a new `Host` structure to include smoothing and support advanced configurations.
- Adjusted training and validation flows to utilize `Host` wrapping the original `ResNet` model.

* chore: clean up outdated comments in ResNet-finetune example

* refactor: reorganize ResNet block APIs and improve method clarity

- Standardized comments in `BasicBlock` and `BottleneckBlock` to use consistent `meta-API` notation.
- Merged duplicate methods (`pinch_factor`, `effective_first_dilation`, `out_planes`) for cleaner implementation.
- Updated method usages to directly reference the respective properties for improved readability and efficiency.

* refactor: rename and standardize Reduction APIs across ResNet blocks

- Replaced `reduction_factor` with `reduce_first` for clarity and consistency.
- Streamlined related methods and comments in `BasicBlock`, `BottleneckBlock`, and `LayerBlock`.
- Added debug print utilities across ResNet modules to help with diagnostics.
- Updated CNA APIs with `map_forward` and adjusted method order for improved readability.
- Enhanced ResNet-finetune example with `--freeze_layers` option and a custom metrics renderer.
- Adjusted default cardinality and pretrained weights handling in ResNet configurations.

* refactor: rename `resnet-finetune` to `resnet_finetune` and update related documentation

- Standardized naming across the example directory, README, and `Cargo.toml`.
- Enhanced pretrained model listing with detailed configurations and descriptions.

* chore: bump version to 0.19.2 across workspace and dependencies

- Updated `Cargo.toml` and `Cargo.lock` files to reflect the new version.

Author: crutcher

Cargo.lock

@@ -10,7 +10,7 @@ members = [
 
 [workspace.package]
 edition = "2024"
-version = "0.19.1"
+version = "0.19.2"
 repository = "https://github.com/crutcher/bimm"
 license = "MIT"
 rust-version = "1.88.0"
@@ -32,7 +32,7 @@ dirs = "^6.0.0"
 # > cudarc = { workspace = true, features = ["nvrtc"], optional = true }
 cudarc = "0.16.6"
 
-bimm-contracts = "^0.19.0"
+bimm-contracts = "^0.19.1"
 
 # Burn coupled-dependencies
 globwalk = "^0.9.1"

Cargo.toml

@@ -75,6 +75,33 @@ let model: ResNet<Wgpu> = prefab
     .with_stochastic_path_depth(0.1);
 ```
 
+### Example [resnet_finetune](examples/resnet_finetune) - Pretrained ResNet finetuning example.
+
+```terminaloutput
+$ cargo run --release -p resnet_finetune -- --pretrained list
+Available pretrained models:
+* "resnet18"
+ResNetContractConfig { layers: [2, 2, 2, 2], num_classes: 1000, stem_width: 64, output_stride: 32, bottleneck_policy: None, normalization: Batch(BatchNormConfig { num_features: 0, epsilon: 1e-5, momentum: 0.1 }), activation: Relu }
+  - "resnet18.tv_in1k": TorchVision ResNet-18
+  - "resnet18.a1_in1k": RSB Paper ResNet-18 a1
+  - "resnet18.a2_in1k": RSB Paper ResNet-18 a2
+  - "resnet18.a3_in1k": RSB Paper ResNet-18 a3
+* "resnet34"
+ResNetContractConfig { layers: [3, 4, 6, 3], num_classes: 1000, stem_width: 64, output_stride: 32, bottleneck_policy: None, normalization: Batch(BatchNormConfig { num_features: 0, epsilon: 1e-5, momentum: 0.1 }), activation: Relu }
+  - "resnet34.tv_in1k": TorchVision ResNet-34
+  - "resnet34.a1_in1k": RSB Paper ResNet-32 a1
+  - "resnet34.a2_in1k": RSB Paper ResNet-32 a2
+  - "resnet34.a3_in1k": RSB Paper ResNet-32 a3
+  - "resnet34.bt_in1k": ResNet-34 pretrained on ImageNet
+* "resnet50"
+ResNetContractConfig { layers: [3, 4, 6, 3], num_classes: 1000, stem_width: 64, output_stride: 32, bottleneck_policy: Some(BottleneckPolicyConfig { pinch_factor: 4 }), normalization: Batch(BatchNormConfig { num_features: 0, epsilon: 1e-5, momentum: 0.1 }), activation: Relu }
+  - "resnet50.tv_in1k": TorchVision ResNet-50
+* "resnet101"
+ResNetContractConfig { layers: [3, 4, 23, 3], num_classes: 1000, stem_width: 64, output_stride: 32, bottleneck_policy: Some(BottleneckPolicyConfig { pinch_factor: 4 }), normalization: Batch(BatchNormConfig { num_features: 0, epsilon: 1e-5, momentum: 0.1 }), activation: Relu }
+  - "resnet101.tv_in1k": TorchVision ResNet-101
+  - "resnet101.a1_in1k": ResNet-101 pretrained on ImageNet
+```
+
 ### [bimm-contracts](https://github.com/crutcher/bimm-contracts) - a crate for static shape contracts for tensors.
 
 [![Crates.io Version](https://img.shields.io/crates/v/bimm-contracts)](https://crates.io/crates/bimm-contracts)

README.md

@@ -12,7 +12,7 @@ description = "bimm-firehose image processing support"
 workspace = true
 
 [dependencies]
-bimm-firehose = { version = "0.19.1", path = "../bimm-firehose" }
+bimm-firehose = { version = "0.19.2", path = "../bimm-firehose" }
 
 burn = { workspace = true, features = ["dataset", "vision", "ndarray"] }
 serde = { workspace = true, features = ["derive"] }

crates/bimm-firehose-image/Cargo.toml

@@ -21,12 +21,13 @@ burn-import = { workspace = true }
 
 crc = { workspace = true }
 serde = { workspace = true, features = ["derive"] }
+serde_json = { workspace = true }
 
 anyhow = { workspace = true }
 num-traits = { workspace = true }
 
 [dev-dependencies]
-burn = { workspace = true, features = ["autodiff"] }
+burn = { workspace = true, features = ["autodiff", "wgpu"] }
 hamcrest = { workspace = true }
 indoc = { workspace = true }
 criterion = { workspace = true }

crates/bimm/Cargo.toml

@@ -57,12 +57,12 @@ pub trait CNA2dMeta {
     /// Number of input channels.
     fn in_channels(&self) -> usize;
 
-    /// Number of groups.
-    fn groups(&self) -> usize;
-
     /// Number of output channels.
     fn out_channels(&self) -> usize;
 
+    /// Number of groups.
+    fn groups(&self) -> usize;
+
     /// Get the stride.
     fn stride(&self) -> [usize; 2];
 }
@@ -88,14 +88,14 @@ impl CNA2dMeta for CNA2dConfig {
         self.conv.channels[0]
     }
 
-    fn groups(&self) -> usize {
-        self.conv.groups
-    }
-
     fn out_channels(&self) -> usize {
         self.conv.channels[1]
     }
 
+    fn groups(&self) -> usize {
+        self.conv.groups
+    }
+
     fn stride(&self) -> [usize; 2] {
         self.conv.stride
     }
@@ -153,15 +153,15 @@ pub struct CNA2d<B: Backend> {
 
 impl<B: Backend> CNA2dMeta for CNA2d<B> {
     fn in_channels(&self) -> usize {
-        self.conv.weight.shape().dims[1] * self.groups()
+        self.conv.weight.dims()[1] * self.groups()
     }
 
-    fn groups(&self) -> usize {
-        self.conv.groups
+    fn out_channels(&self) -> usize {
+        self.conv.weight.dims()[0]
     }
 
-    fn out_channels(&self) -> usize {
-        self.conv.weight.shape().dims[0]
+    fn groups(&self) -> usize {
+        self.conv.groups
     }
 
     fn stride(&self) -> [usize; 2] {
@@ -193,17 +193,17 @@ impl<B: Backend> CNA2d<B> {
         &self,
         input: Tensor<B, 4>,
     ) -> Tensor<B, 4> {
-        self.hook_forward(input, |x| x)
+        self.map_forward(input, |x| x)
     }
 
-    /// Hooked Forward Pass.
+    /// Mapping Forward Pass.
     ///
-    /// Applies the hook after normalization but before activation.
+    /// Applies the callback fn after normalization but before activation.
     ///
     /// ```rust,ignore
     /// let x = self.conv.forward(input);
     /// let x = self.norm.forward(x);
-    /// let x = hook(x);
+    /// let x = f(x);
     /// let x = self.act.forward(x);
     /// return x
     /// ```
@@ -212,14 +212,15 @@ impl<B: Backend> CNA2d<B> {
     ///
     /// - `input`: \
     ///   ``[batch, in_channels, in_height=out_height*stride, in_width=out_width*stride]``.
+    /// - `f`: a callback endofunction, from/to ``[batch, in_channels, out_height, out_width]``.
     ///
     /// # Returns
     ///
     /// ``[batch, out_channels, out_height, out_width]``
-    pub fn hook_forward<F>(
+    pub fn map_forward<F>(
         &self,
         input: Tensor<B, 4>,
-        hook: F,
+        f: F,
     ) -> Tensor<B, 4>
     where
         F: FnOnce(Tensor<B, 4>) -> Tensor<B, 4>,
@@ -254,7 +255,7 @@ impl<B: Backend> CNA2d<B> {
 
         let x = self.norm.forward(x);
 
-        let x = hook(x);
+        let x = f(x);
 
         let x = self.act.forward(x);
 
@@ -345,7 +346,7 @@ mod tests {
         {
             let hook = |x| x * 2.0;
 
-            let output = layer.hook_forward(input.clone(), hook);
+            let output = layer.map_forward(input.clone(), hook);
             let expected = {
                 let x = layer.conv.forward(input.clone());
                 let x = layer.norm.forward(x);