Experiment with resnet (#1128)

* Add some preliminary support for resnet.

* Add an actual resnet example.

candle-transformers/src/models/mod.rs

@@ -14,6 +14,7 @@ pub mod quantized_mixformer;
 pub mod quantized_mpt;
 pub mod quantized_stable_lm;
 pub mod quantized_t5;
+pub mod resnet;
 pub mod segment_anything;
 pub mod stable_diffusion;
 pub mod stable_lm;

candle-transformers/src/models/resnet.rs

@@ -0,0 +1,131 @@
+//! ResNet implementation.
+//!
+//! See "Deep Residual Learning for Image Recognition" He et al. 2015
+//! <https://arxiv.org/abs/1512.03385>
+use candle::{Result, D};
+use candle_nn::{batch_norm, Conv2d, Func, VarBuilder};
+
+fn conv2d(
+    c_in: usize,
+    c_out: usize,
+    ksize: usize,
+    padding: usize,
+    stride: usize,
+    vb: VarBuilder,
+) -> Result<Conv2d> {
+    let conv2d_cfg = candle_nn::Conv2dConfig {
+        stride,
+        padding,
+        ..Default::default()
+    };
+    candle_nn::conv2d_no_bias(c_in, c_out, ksize, conv2d_cfg, vb)
+}
+
+fn downsample(c_in: usize, c_out: usize, stride: usize, vb: VarBuilder) -> Result<Func> {
+    if stride != 1 || c_in != c_out {
+        let conv = conv2d(c_in, c_out, 1, 0, stride, vb.pp(0))?;
+        let bn = batch_norm(c_out, 1e-5, vb.pp(1))?;
+        Ok(Func::new(move |xs| xs.apply(&conv)?.apply(&bn)))
+    } else {
+        Ok(Func::new(|xs| Ok(xs.clone())))
+    }
+}
+
+fn basic_block(c_in: usize, c_out: usize, stride: usize, vb: VarBuilder) -> Result<Func> {
+    let conv1 = conv2d(c_in, c_out, 3, 1, stride, vb.pp("conv1"))?;
+    let bn1 = batch_norm(c_out, 1e-5, vb.pp("bn1"))?;
+    let conv2 = conv2d(c_out, c_out, 3, 1, 1, vb.pp("conv2"))?;
+    let bn2 = batch_norm(c_out, 1e-5, vb.pp("bn2"))?;
+    let downsample = downsample(c_in, c_out, stride, vb.pp("downsample"))?;
+    Ok(Func::new(move |xs| {
+        let ys = xs
+            .apply(&conv1)?
+            .apply(&bn1)?
+            .relu()?
+            .apply(&conv2)?
+            .apply(&bn2)?;
+        (xs.apply(&downsample)? + ys)?.relu()
+    }))
+}
+
+fn basic_layer(
+    c_in: usize,
+    c_out: usize,
+    stride: usize,
+    cnt: usize,
+    vb: VarBuilder,
+) -> Result<Func> {
+    let mut layers = Vec::with_capacity(cnt);
+    for index in 0..cnt {
+        let l_in = if index == 0 { c_in } else { c_out };
+        let stride = if index == 0 { stride } else { 1 };
+        layers.push(basic_block(l_in, c_out, stride, vb.pp(index))?)
+    }
+    Ok(Func::new(move |xs| {
+        let mut xs = xs.clone();
+        for layer in layers.iter() {
+            xs = xs.apply(layer)?
+        }
+        Ok(xs)
+    }))
+}
+
+fn resnet(
+    nclasses: Option<usize>,
+    c1: usize,
+    c2: usize,
+    c3: usize,
+    c4: usize,
+    vb: VarBuilder,
+) -> Result<Func> {
+    let conv1 = conv2d(3, 64, 7, 3, 2, vb.pp("conv1"))?;
+    let bn1 = batch_norm(64, 1e-5, vb.pp("bn1"))?;
+    let layer1 = basic_layer(64, 64, 1, c1, vb.pp("layer1"))?;
+    let layer2 = basic_layer(64, 128, 2, c2, vb.pp("layer2"))?;
+    let layer3 = basic_layer(128, 256, 2, c3, vb.pp("layer3"))?;
+    let layer4 = basic_layer(256, 512, 2, c4, vb.pp("layer4"))?;
+    let fc = match nclasses {
+        None => None,
+        Some(nclasses) => {
+            let linear = candle_nn::linear(512, nclasses, vb.pp("fc"))?;
+            Some(linear)
+        }
+    };
+    Ok(Func::new(move |xs| {
+        let xs = xs
+            .apply(&conv1)?
+            .apply(&bn1)?
+            .relu()?
+            .pad_with_same(D::Minus1, 1, 1)?
+            .pad_with_same(D::Minus2, 1, 1)?
+            .max_pool2d_with_stride(3, 2)?
+            .apply(&layer1)?
+            .apply(&layer2)?
+            .apply(&layer3)?
+            .apply(&layer4)?
+            .mean(D::Minus1)?
+            .mean(D::Minus1)?;
+        match &fc {
+            None => Ok(xs),
+            Some(fc) => xs.apply(fc),
+        }
+    }))
+}
+
+/// Creates a ResNet-18 model.
+pub fn resnet18(num_classes: usize, vb: VarBuilder) -> Result<Func> {
+    resnet(Some(num_classes), 2, 2, 2, 2, vb)
+}
+
+pub fn resnet18_no_final_layer(vb: VarBuilder) -> Result<Func> {
+    resnet(None, 2, 2, 2, 2, vb)
+}
+
+/// Creates a ResNet-34 model.
+pub fn resnet34(num_classes: usize, vb: VarBuilder) -> Result<Func> {
+    resnet(Some(num_classes), 3, 4, 6, 3, vb)
+}
+
+pub fn resnet34_no_final_layer(vb: VarBuilder) -> Result<Func> {
+    resnet(None, 3, 4, 6, 3, vb)
+}