docs: update ResNet examples for clarity and improved prefab usage

- Revised ResNet examples in documentation to align with updated prefab and weight-fetching methods.
- Enhanced example readability with consistent formatting and commented steps for better guidance.
- Updated references to new prefab and weight-loading APIs.

Author: crutcher

README.md

@@ -45,28 +45,30 @@ See the [CONTRIBUTING](CONTRIBUTING.md) guide for build and contribution instruc
 
 #### Example
 
-Example of building a pretrained ResNet-18 module:
+Example of building a pretrained model:
 
 ```rust,no_run
-use bimm::cache::fetch_model_weights;
-use bimm::models::resnet::{ResNet, ResNetAbstractConfig};
-use burn::backend::NdArray;
+use burn::backend::Wgpu;
+use bimm::cache::disk::DiskCacheConfig;
+use bimm::models::resnet::{PREFAB_RESNET_MAP, ResNet};
 
 let device = Default::default();
 
-let source =
-    "https://download.pytorch.org/models/resnet18-f37072fd.pth";
-let source_classes = 1000;
-let weights_path= fetch_model_weights(source).unwrap();
+let prefab = PREFAB_RESNET_MAP.expect_lookup_prefab("resnet18");
 
-let my_classes = 10;
+let weights = prefab
+    .expect_lookup_pretrained_weights("tv_in1k")
+    .fetch_weights(&DiskCacheConfig::default())
+    .expect("Failed to fetch weights");
 
-let model: ResNet<NdArray> = ResNetAbstractConfig::resnet18(source_classes)
+let model: ResNet<Wgpu> = prefab
+    .to_config()
     .to_structure()
     .init(&device)
-    .load_pytorch_weights(weights_path)
-    .expect("Model should be loaded successfully")
-    .with_classes(my_classes)
+    .load_pytorch_weights(weights)
+    .expect("Failed to load weights")
+    // re-head the model to 10 classes:
+    .with_classes(10)
     // Enable (drop_block_prob) stochastic block drops for training:
     .with_stochastic_drop_block(0.2)
     // Enable (drop_path_prob) stochastic depth for training:

crates/bimm/README.md

@@ -5,29 +5,30 @@
 
 This is a Rust crate for image models, inspired by the Python `timm` package.
 
-Examples of loading pre-trained ResNet-18 model:
+Examples of loading pretrained model:
 
 ```rust,no_run
-use bimm::cache::fetch_model_weights;
-use bimm::models::resnet::{ResNet, ResNetAbstractConfig};
 use burn::backend::Wgpu;
+use bimm::cache::disk::DiskCacheConfig;
+use bimm::models::resnet::{PREFAB_RESNET_MAP, ResNet};
 
-type B = Wgpu;
 let device = Default::default();
 
-let source =
-    "https://download.pytorch.org/models/resnet18-f37072fd.pth";
-let source_classes = 1000;
-let weights_path= fetch_model_weights(source).unwrap();
+let prefab = PREFAB_RESNET_MAP.expect_lookup_prefab("resnet18");
 
-let my_classes = 10;
+let weights = prefab
+    .expect_lookup_pretrained_weights("tv_in1k")
+    .fetch_weights(&DiskCacheConfig::default())
+    .expect("Failed to fetch weights");
 
-let model: ResNet<B> = ResNetAbstractConfig::resnet18(source_classes)
+let model: ResNet<Wgpu> = prefab
+    .to_config()
     .to_structure()
     .init(&device)
-    .load_pytorch_weights(weights_path)
-    .expect("Model should be loaded successfully")
-    .with_classes(my_classes)
+    .load_pytorch_weights(weights)
+    .expect("Failed to load weights")
+    // re-head the model to 10 classes:
+    .with_classes(10)
     // Enable (drop_block_prob) stochastic block drops for training:
     .with_stochastic_drop_block(0.2)
     // Enable (drop_path_prob) stochastic depth for training:

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

@@ -5,6 +5,8 @@
 //!
 //! ## Example
 //!
+//! Examples of loading pretrained model:
+//!
 //! ```rust,no_run
 //! use burn::backend::NdArray;
 //! use bimm::cache::disk::DiskCacheConfig;
@@ -25,8 +27,12 @@
 //!     .init(&device)
 //!     .load_pytorch_weights(weights)
 //!     .expect("Failed to load weights")
+//!     // re-head the model to 10 classes:
 //!     .with_classes(10)
-//!     .with_stochastic_drop_block(0.2);
+//!     // Enable (drop_block_prob) stochastic block drops for training:
+//!     .with_stochastic_drop_block(0.2)
+//!     // Enable (drop_path_prob) stochastic depth for training:
+//!     .with_stochastic_path_depth(0.1);
 //! ```
 //!
 //! ## Configuration