FastViT fixes. (#2452)

* correct optional SE layer dimensions.
 * head_dim instead of num_heads is 32.
 * update test example output.

README.md

@@ -241,7 +241,7 @@ If you have an addition to this list, please submit a pull request.
         - Parler-TTS, text-to-speech model.
     - Computer Vision Models.
         - DINOv2, ConvMixer, EfficientNet, ResNet, ViT, VGG, RepVGG, ConvNeXT,
-          ConvNeXTv2, MobileOne, EfficientVit (MSRA), MobileNetv4, Hiera.
+          ConvNeXTv2, MobileOne, EfficientVit (MSRA), MobileNetv4, Hiera, FastViT.
         - yolo-v3, yolo-v8.
         - Segment-Anything Model (SAM).
         - SegFormer.

candle-examples/examples/fastvit/README.md

@@ -12,9 +12,9 @@ $ cargo run --example fastvit --release -- --image candle-examples/examples/yolo
 
 loaded image Tensor[dims 3, 256, 256; f32]
 model built
-mountain bike, all-terrain bike, off-roader: 43.45%
-bicycle-built-for-two, tandem bicycle, tandem: 14.16%
-unicycle, monocycle     : 4.12%
-crash helmet            : 2.26%
-alp                     : 1.40%
+mountain bike, all-terrain bike, off-roader: 52.67%
+bicycle-built-for-two, tandem bicycle, tandem: 7.93%
+unicycle, monocycle     : 3.46%
+maillot                 : 1.32%
+crash helmet            : 1.28%
 ```

candle-transformers/src/models/fastvit.rs

@@ -339,8 +339,8 @@ fn positional_encoding(dim: usize, vb: VarBuilder) -> Result<Func<'static>> {
 fn attention(dim: usize, vb: VarBuilder) -> Result<Func<'static>> {
     let qkv = linear_no_bias(dim, dim * 3, vb.pp("qkv"))?;
     let proj = linear(dim, dim, vb.pp("proj"))?;
-    let num_heads = 32;
-    let head_dim = dim / num_heads;
+    let head_dim = 32;
+    let num_heads = dim / head_dim;
     let scale = (head_dim as f64).powf(-0.5);
 
     Ok(Func::new(move |xs| {
@@ -434,7 +434,7 @@ fn fastvit_patch_embed(
 ) -> Result<Func<'static>> {
     let lk = conv_norm(in_channels, out_channels, 7, 2, vb.pp("proj.0.large_conv"))?;
     let sk = conv_norm(in_channels, out_channels, 3, 2, vb.pp("proj.0.small_conv"))?;
-    let se = squeeze_and_excitation(out_channels, out_channels / 16, vb.pp("proj.0.se"));
+    let se = squeeze_and_excitation(out_channels, out_channels / 4, vb.pp("proj.0.se"));
     let mb = mobileone_block(out_channels, out_channels, 1, 1, 0, true, vb.pp("proj.1"))?;
 
     Ok(Func::new(move |xs| {