feat: introduce multi-backend support and update configurations

- Added support for `WGPU`, `CUDA`, and `Metal` backends in examples (`resnet_tiny`, `resnet-finetune`, and `swin_tiny`).
- Refactored backend initialization logic to dynamically handle multiple features.
- Updated `Cargo.toml` files to add feature flags (`wgpu`, `cuda`, `metal`) for enhanced flexibility.
- Modified `.run` configurations to include appropriate backend-specific settings.
- Upgraded `aho-corasick` dependency from `1.1.3` to `1.1.4` in `Cargo.lock`.

Author: crutcher

.run/Fix.run.xml

@@ -1,7 +1,7 @@
 <component name="ProjectRunConfigurationManager">
   <configuration default="false" name="Fix" type="CargoCommandRunConfiguration" factoryName="Cargo Command" focusToolWindowBeforeRun="true">
     <option name="buildProfileId" value="test" />
-    <option name="command" value="clippy --fix --allow-dirty --allow-staged" />
+    <option name="command" value="clippy --fix --allow-dirty --allow-staged --features wgpu" />
     <option name="workingDirectory" value="file://$PROJECT_DIR$" />
     <envs>
       <env name="CINIC10_PATH" value="$CINIC10_PATH$" />

.run/Test.run.xml

@@ -1,7 +1,7 @@
 <component name="ProjectRunConfigurationManager">
   <configuration default="false" name="Test" type="CargoCommandRunConfiguration" factoryName="Cargo Command" focusToolWindowBeforeRun="true">
     <option name="buildProfileId" value="test" />
-    <option name="command" value="test" />
+    <option name="command" value="test --features wgpu" />
     <option name="workingDirectory" value="file://$PROJECT_DIR$" />
     <envs>
       <env name="CINIC10_PATH" value="/media/Data/CINIC-10" />

Cargo.lock

@@ -10,7 +10,7 @@ members = [
 
 [workspace.package]
 edition = "2024"
-version = "0.19.0"
+version = "0.19.1"
 repository = "https://github.com/crutcher/bimm"
 license = "MIT"
 rust-version = "1.88.0"

Cargo.toml

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

crates/bimm-firehose-image/Cargo.toml

@@ -411,7 +411,7 @@ mod tests {
         let output = block.forward(input);
 
         assert_shape_contract!(
-            ["batch", "out_channels", "out_height", "out_width"],
+            ["batch", "out_planes", "out_height", "out_width"],
             &output,
             &[
                 ("batch", batch_size),

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

@@ -11,9 +11,9 @@ workspace = true
 
 [dependencies]
 anyhow = { workspace = true }
-bimm = { version = "0.19.0", path = "../../crates/bimm" }
+bimm = { version = "0.19.1", path = "../../crates/bimm" }
 
-burn = { workspace = true, features = ["network", "train", "autodiff", "vision", "cuda"] }
+burn = { workspace = true, features = ["network", "train", "autodiff", "vision"] }
 
 clap = { workspace = true, features = ["derive"] }
 
@@ -24,7 +24,9 @@ rand = { workspace = true }
 serde = { workspace = true, features = ["std", "derive"] }
 tar = { workspace = true }
 
-
 [features]
 default = []
+wgpu = ["burn/wgpu"]
+cuda = ["burn/cuda"]
+metal = ["burn/metal"]
 

examples/resnet-finetune/Cargo.toml

@@ -9,7 +9,7 @@ use crate::data::{ClassificationBatch, ClassificationBatcher};
 use crate::dataset::{CLASSES, PlanetLoader, download};
 use bimm::cache::disk::DiskCacheConfig;
 use bimm::models::resnet::{PREFAB_RESNET_MAP, ResNet, ResNetContractConfig};
-use burn::backend::{Autodiff, Cuda};
+use burn::backend::Autodiff;
 use burn::config::Config;
 use burn::data::dataloader::DataLoaderBuilder;
 use burn::data::dataset::transform::ShuffledDataset;
@@ -150,21 +150,26 @@ fn main() -> anyhow::Result<()> {
 
     let _source_tree = download();
 
-    let device = Default::default();
-    train::<Autodiff<Cuda>>(&args, &device)
+    #[cfg(feature = "wgpu")]
+    return train::<Autodiff<burn::backend::Wgpu>>(&args);
+
+    #[cfg(feature = "cuda")]
+    return train::<Autodiff<burn::backend::Cuda>>(&args);
+
+    #[cfg(feature = "metal")]
+    return train::<Autodiff<burn::backend::Metal>>(&args);
 }
 
 #[must_use]
-pub fn train<B: AutodiffBackend>(
-    args: &Args,
-    device: &B::Device,
-) -> anyhow::Result<()> {
+pub fn train<B: AutodiffBackend>(args: &Args) -> anyhow::Result<()> {
+    let device: B::Device = Default::default();
+
     // 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).expect("Failed to create artifacts directory");
 
-    B::seed(device, args.seed);
+    B::seed(&device, args.seed);
 
     let disk_cache = DiskCacheConfig::default();
 
@@ -180,7 +185,7 @@ pub fn train<B: AutodiffBackend>(
     let model: ResNet<B> = resnet_config
         .clone()
         .to_structure()
-        .init(device)
+        .init(&device)
         .load_pytorch_weights(weights)
         .expect("Failed to load pretrained weights")
         .with_classes(CLASSES.len())

examples/resnet-finetune/src/main.rs

@@ -10,11 +10,11 @@ repository.workspace = true
 workspace = true
 
 [dependencies]
-bimm = { version = "0.19.0", path = "../../crates/bimm" }
-bimm-firehose = { version = "0.19.0", path = "../../crates/bimm-firehose" }
-bimm-firehose-image = { version = "0.19.0", path = "../../crates/bimm-firehose-image" }
+bimm = { version = "0.19.1", path = "../../crates/bimm" }
+bimm-firehose = { version = "0.19.1", path = "../../crates/bimm-firehose" }
+bimm-firehose-image = { version = "0.19.1", path = "../../crates/bimm-firehose-image" }
 
-burn = { workspace = true, features = ["train", "autodiff", "cuda", "fusion", "vision", "dataset"] }
+burn = { workspace = true, features = ["train", "autodiff", "fusion", "vision", "dataset"] }
 serde = { workspace = true, features = ["derive"] }
 
 clap = { workspace = true, features = ["derive"] }
@@ -23,4 +23,7 @@ rand = { workspace = true }
 anyhow = { workspace = true }
 
 [features]
-nightly = []
+default = []
+wgpu = ["burn/wgpu"]
+cuda = ["burn/cuda"]
+metal = ["burn/metal"]

examples/resnet_tiny/Cargo.toml

@@ -19,7 +19,7 @@ use bimm_firehose_image::augmentation::orientation::flip::HorizontalFlipStage;
 use bimm_firehose_image::burn_support::{ImageToTensorData, stack_tensor_data_column};
 use bimm_firehose_image::loader::{ImageLoader, ResizeSpec};
 use bimm_firehose_image::{ColorType, ImageShape};
-use burn::backend::{Autodiff, Cuda};
+use burn::backend::Autodiff;
 use burn::data::dataloader::{DataLoaderBuilder, Dataset};
 use burn::data::dataset::transform::ShuffledDataset;
 use burn::lr_scheduler::cosine::CosineAnnealingLrSchedulerConfig;
@@ -126,10 +126,15 @@ pub struct Args {
 
 fn main() -> anyhow::Result<()> {
     let args = Args::parse();
-    type B = Autodiff<Cuda>;
 
-    let device = Default::default();
-    backend_main::<B>(&args, &device)
+    #[cfg(feature = "wgpu")]
+    return backend_main::<Autodiff<burn::backend::Wgpu>>(&args);
+
+    #[cfg(feature = "cuda")]
+    return backend_main::<Autodiff<burn::backend::Cuda>>(&args);
+
+    #[cfg(feature = "metal")]
+    return backend_main::<Autodiff<burn::backend::Metal>>(&args);
 }
 
 /// Create the artifact directory for saving training artifacts.
@@ -140,17 +145,16 @@ fn create_artifact_dir(artifact_dir: &str) {
 }
 
 /// Train the model with the given configuration and devices.
-pub fn backend_main<B: AutodiffBackend>(
-    args: &Args,
-    device: &B::Device,
-) -> anyhow::Result<()> {
+pub fn backend_main<B: AutodiffBackend>(args: &Args) -> anyhow::Result<()> {
+    let device: B::Device = Default::default();
+
     let image_shape = ImageShape {
         height: 32,
         width: 32,
     };
     let num_classes = 10;
 
-    B::seed(device, args.seed);
+    B::seed(&device, args.seed);
 
     let prefab = PREFAB_RESNET_MAP.expect_lookup_prefab(&args.resnet_prefab);
 
@@ -159,7 +163,7 @@ pub fn backend_main<B: AutodiffBackend>(
         .with_activation(ActivationConfig::Gelu)
         .to_structure();
 
-    let resnet: ResNet<B> = resnet_config.init(device);
+    let resnet: ResNet<B> = resnet_config.init(&device);
 
     let resnet: ResNet<B> = match &args.resnet_pretrained {
         Some(pretrained) => {