coreylowman · coreylowman · Feb 24, 2023 · Feb 24, 2023
diff --git a/Cargo.toml b/Cargo.toml
@@ -55,4 +55,16 @@ indicatif = "0.16.2"
 
 [build-dependencies]
 rustc_version = "0.4.0"
-glob = "0.3.0"
+glob = "0.3.0"
+
+[[bench]]
+name = "batchnorm2d"
+harness = false
+
+[[bench]]
+name = "conv2d"
+harness = false
+
+[[bench]]
+name = "sum"
+harness = false
diff --git a/benches/README.md b/benches/README.md
@@ -0,0 +1,10 @@
+# Benchmarks
+
+This folder contains a couple of simple utility scripts for running
+forward and backward passes of core ops.
+
+- `cargo bench --bench batchnorm2d`
+- `cargo bench --bench sum`
+- `cargo +nightly bench --bench conv2d`
+
+Additionally you can pass `-F cuda` to use a Cuda.
diff --git a/benches/batchnorm2d.rs b/benches/batchnorm2d.rs
@@ -0,0 +1,38 @@
+use std::time::Instant;
+
+use dfdx::prelude::*;
+
+#[cfg(feature = "cuda")]
+type Dev = Cuda;
+
+#[cfg(not(feature = "cuda"))]
+type Dev = Cpu;
+
+type Model = BatchNorm2D<512>;
+type Dtype = f32;
+type InputShape = Rank4<64, 512, 28, 28>;
+
+fn main() {
+    println!("Benchmarking `BatchNorm2D`");
+    println!("Device {}", std::any::type_name::<Dev>());
+    println!("Dtype {}", std::any::type_name::<Dtype>());
+    println!("Input shape {}", std::any::type_name::<InputShape>());
+    println!();
+
+    let dev: Dev = Default::default();
+    let mut m = dev.build_module::<Model, Dtype>();
+
+    loop {
+        let img: Tensor<InputShape, Dtype, _> = dev.sample_normal();
+
+        let start = Instant::now();
+        let out = m.forward_mut(img.traced());
+        let loss = out.square().mean();
+        let fwd_dur = start.elapsed();
+
+        let start = Instant::now();
+        let _ = loss.backward();
+        let bwd_dur = start.elapsed();
+        println!("fwd={:?} bwd={:?}", fwd_dur, bwd_dur);
+    }
+}
diff --git a/benches/conv2d.rs b/benches/conv2d.rs
@@ -0,0 +1,46 @@
+#![cfg_attr(feature = "nightly", feature(generic_const_exprs))]
+
+#[cfg(feature = "nightly")]
+fn main() {
+    use std::time::Instant;
+
+    use dfdx::prelude::*;
+
+    #[cfg(feature = "cuda")]
+    type Dev = Cuda;
+
+    #[cfg(not(feature = "cuda"))]
+    type Dev = Cpu;
+
+    type Model = Conv2D<128, 256, 4>;
+    type Dtype = f32;
+    type InputShape = Rank4<64, 128, 28, 28>;
+
+    println!("Benchmarking `Conv2D`");
+    println!("Device {}", std::any::type_name::<Dev>());
+    println!("Dtype {}", std::any::type_name::<Dtype>());
+    println!("Input shape {}", std::any::type_name::<InputShape>());
+    println!();
+
+    let dev: Dev = Default::default();
+    let mut m = dev.build_module::<Model, Dtype>();
+
+    loop {
+        let img: Tensor<InputShape, Dtype, _> = dev.sample_normal();
+
+        let start = Instant::now();
+        let out = m.forward_mut(img.traced());
+        let loss = out.square().mean();
+        let fwd_dur = start.elapsed();
+
+        let start = Instant::now();
+        let _ = loss.backward();
+        let bwd_dur = start.elapsed();
+        println!("fwd={:?} bwd={:?}", fwd_dur, bwd_dur);
+    }
+}
+
+#[cfg(not(feature = "nightly"))]
+fn main() {
+    panic!("Run with `cargo +nightly run ...` to run this example.");
+}
diff --git a/benches/sum.rs b/benches/sum.rs
@@ -0,0 +1,35 @@
+use std::time::Instant;
+
+use dfdx::prelude::*;
+
+#[cfg(feature = "cuda")]
+type Dev = Cuda;
+
+#[cfg(not(feature = "cuda"))]
+type Dev = Cpu;
+
+type Dtype = f32;
+type InputShape = Rank4<32, 64, 128, 256>;
+
+fn main() {
+    println!("Benchmarking `sum`");
+    println!("Device {}", std::any::type_name::<Dev>());
+    println!("Dtype {}", std::any::type_name::<Dtype>());
+    println!("Input shape {}", std::any::type_name::<InputShape>());
+    println!();
+
+    let dev: Dev = Default::default();
+
+    loop {
+        let img: Tensor<InputShape, Dtype, _> = dev.sample_normal();
+
+        let start = Instant::now();
+        let y = img.traced().sum();
+        let fwd_dur = start.elapsed();
+
+        let start = Instant::now();
+        let _ = y.backward();
+        let bwd_dur = start.elapsed();
+        println!("fwd={:?} bwd={:?}", fwd_dur, bwd_dur);
+    }
+}