#![cfg_attr(feature = "nightly", feature(generic_const_exprs))]
#[cfg(feature = "nightly")]
fn main() {
use dfdx::prelude::*;
use std::time::Instant;
type BasicBlock<const C: usize> = Residual<(
Conv2D<C, C, 3, 1, 1>,
BatchNorm2D<C>,
ReLU,
Conv2D<C, C, 3, 1, 1>,
BatchNorm2D<C>,
)>;
type Downsample<const C: usize, const D: usize> = GeneralizedResidual<
(
Conv2D<C, D, 3, 2, 1>,
BatchNorm2D<D>,
ReLU,
Conv2D<D, D, 3, 1, 1>,
BatchNorm2D<D>,
),
(Conv2D<C, D, 1, 2, 0>, BatchNorm2D<D>),
>;
type Head = (
Conv2D<3, 64, 7, 2, 3>,
BatchNorm2D<64>,
ReLU,
MaxPool2D<3, 2, 1>,
);
type Resnet18<const NUM_CLASSES: usize> = (
Head,
(BasicBlock<64>, ReLU, BasicBlock<64>, ReLU),
(Downsample<64, 128>, ReLU, BasicBlock<128>, ReLU),
(Downsample<128, 256>, ReLU, BasicBlock<256>, ReLU),
(Downsample<256, 512>, ReLU, BasicBlock<512>, ReLU),
(AvgPoolGlobal, Linear<512, NUM_CLASSES>),
);
let dev = AutoDevice::default();
let m = dev.build_module::<Resnet18<1000>, f32>();
let x: Tensor<Rank3<3, 224, 224>, f32, _> = dev.sample_normal();
const PROBES: u32 = 10;
let start = Instant::now();
for _ in 0..PROBES {
let _y = m.forward(x.clone());
}
println!("Average unbatched forward: {:?}", start.elapsed() / PROBES);
let x: Tensor<Rank4<16, 3, 224, 224>, f32, _> = dev.sample_normal();
let start = Instant::now();
for _ in 0..PROBES {
let _y = m.forward(x.clone());
}
println!(
"Average batched (16) forward: {:?}",
start.elapsed() / PROBES
);
}
#[cfg(not(feature = "nightly"))]
fn main() {
panic!("Run with `cargo +nightly run ...` to run this example.");
}