diff --git a/candle-examples/examples/beit/README.md b/candle-examples/examples/beit/README.md new file mode 100644 index 0000000000..23af1e32f5 --- /dev/null +++ b/candle-examples/examples/beit/README.md @@ -0,0 +1,20 @@ +# candle-beit + +[Beit](https://arxiv.org/abs/2106.08254) is a computer vision model. +In this example, it is used as an ImageNet classifier: the model returns the +probability for the image to belong to each of the 1000 ImageNet categories. + +## Running some example + +```bash +cargo run --example beit --release -- --image candle-examples/examples/yolo-v8/assets/bike.jpg + +> mountain bike, all-terrain bike, off-roader: 56.16% +> bicycle-built-for-two, tandem bicycle, tandem: 3.08% +> maillot : 2.23% +> alp : 0.88% +> crash helmet : 0.85% + +``` + +![Leading group, Giro d'Italia 2021](../yolo-v8/assets/bike.jpg) diff --git a/candle-examples/examples/beit/main.rs b/candle-examples/examples/beit/main.rs new file mode 100644 index 0000000000..5ef2a6ae0c --- /dev/null +++ b/candle-examples/examples/beit/main.rs @@ -0,0 +1,79 @@ +//! BEiT: BERT Pre-Training of Image Transformers +//! https://github.com/microsoft/unilm/tree/master/beit + +#[cfg(feature = "mkl")] +extern crate intel_mkl_src; + +#[cfg(feature = "accelerate")] +extern crate accelerate_src; + +use clap::Parser; + +use candle::{DType, Device, IndexOp, Result, Tensor, D}; +use candle_nn::{Module, VarBuilder}; +use candle_transformers::models::beit; + +/// Loads an image from disk using the image crate, this returns a tensor with shape +/// (3, 384, 384). Beit special normalization is applied. +pub fn load_image384_beit_norm>(p: P) -> Result { + let img = image::io::Reader::open(p)? + .decode() + .map_err(candle::Error::wrap)? + .resize_to_fill(384, 384, image::imageops::FilterType::Triangle); + let img = img.to_rgb8(); + let data = img.into_raw(); + let data = Tensor::from_vec(data, (384, 384, 3), &Device::Cpu)?.permute((2, 0, 1))?; + let mean = Tensor::new(&[0.5f32, 0.5, 0.5], &Device::Cpu)?.reshape((3, 1, 1))?; + let std = Tensor::new(&[0.5f32, 0.5, 0.5], &Device::Cpu)?.reshape((3, 1, 1))?; + (data.to_dtype(candle::DType::F32)? / 255.)? + .broadcast_sub(&mean)? + .broadcast_div(&std) +} + +#[derive(Parser)] +struct Args { + #[arg(long)] + model: Option, + + #[arg(long)] + image: String, + + /// Run on CPU rather than on GPU. + #[arg(long)] + cpu: bool, +} + +pub fn main() -> anyhow::Result<()> { + let args = Args::parse(); + + let device = candle_examples::device(args.cpu)?; + + let image = load_image384_beit_norm(args.image)?.to_device(&device)?; + println!("loaded image {image:?}"); + + let model_file = match args.model { + None => { + let api = hf_hub::api::sync::Api::new()?; + let api = api.model("vincent-espitalier/candle-beit".into()); + api.get("beit_base_patch16_384.in22k_ft_in22k_in1k_adapted.safetensors")? + } + Some(model) => model.into(), + }; + let vb = unsafe { VarBuilder::from_mmaped_safetensors(&[model_file], DType::F32, &device)? }; + let model = beit::vit_base(vb)?; + println!("model built"); + let logits = model.forward(&image.unsqueeze(0)?)?; + let prs = candle_nn::ops::softmax(&logits, D::Minus1)? + .i(0)? + .to_vec1::()?; + let mut prs = prs.iter().enumerate().collect::>(); + prs.sort_by(|(_, p1), (_, p2)| p2.total_cmp(p1)); + for &(category_idx, pr) in prs.iter().take(5) { + println!( + "{:24}: {:.2}%", + candle_examples::imagenet::CLASSES[category_idx], + 100. * pr + ); + } + Ok(()) +} diff --git a/candle-transformers/src/models/beit.rs b/candle-transformers/src/models/beit.rs new file mode 100644 index 0000000000..c534032cc2 --- /dev/null +++ b/candle-transformers/src/models/beit.rs @@ -0,0 +1,367 @@ +use candle::{DType, IndexOp, Result, Tensor, D}; +use candle_nn::{layer_norm, LayerNorm, Linear, Module, VarBuilder}; + +const IMG_SIZE: usize = 384; +const PATCH_SIZE: usize = 16; +const NUM_CLASSES: usize = 1000; +const WINDOW_SIZE: usize = IMG_SIZE / PATCH_SIZE; // 384 / 16 = 24 +const NB_TOKENS: usize = WINDOW_SIZE * WINDOW_SIZE + 1; // 24 * 24 + 1 = 577 + +fn linear(vb: VarBuilder, in_dim: usize, out_dim: usize, bias: bool) -> Result { + if bias { + candle_nn::linear(in_dim, out_dim, vb) + } else { + candle_nn::linear_no_bias(in_dim, out_dim, vb) + } +} + +#[derive(Debug)] +struct Attention { + qkv: Linear, + proj: Linear, + relative_position_bias_table: Tensor, + relative_position_index: Tensor, + num_heads: usize, + scale: f64, +} + +impl Attention { + fn new( + vb: VarBuilder, + dim: usize, + num_heads: usize, + qkv_bias: bool, + proj_bias: bool, + relative_position_index: &Tensor, + ) -> Result { + let qkv = linear(vb.pp("qkv"), dim, dim * 3, qkv_bias)?; + let proj = linear(vb.pp("proj"), dim, dim, proj_bias)?; + // num_relative_distance = token-token(47x47) + token-CLS(1) + CLS-token(1) + CLS-CLS(1) = 2212 + let num_relative_distance = (2 * WINDOW_SIZE - 1) * (2 * WINDOW_SIZE - 1) + 3; + let relative_position_bias_table = vb.get( + (num_relative_distance, num_heads), + "relative_position_bias_table", + )?; + let relative_position_index = relative_position_index.clone(); + let scale = 1. / ((dim / num_heads) as f64).sqrt(); + Ok(Self { + qkv, + proj, + relative_position_bias_table, + relative_position_index, + num_heads, + scale, + }) + } +} + +impl Attention { + fn _get_rel_pos_bias(&self) -> Result { + self.relative_position_bias_table + .index_select( + &self + .relative_position_index + .flatten_all()? + .to_dtype(DType::U32)?, + 0, + )? + .reshape((NB_TOKENS, NB_TOKENS, ()))? + .transpose(0, 1)? // 102 + .transpose(0, 2)? // 201 + .contiguous()? + .unsqueeze(0) + } +} + +impl Module for Attention { + fn forward(&self, xs: &Tensor) -> Result { + let (b, n, c) = xs.dims3()?; + let qkv = self + .qkv + .forward(xs)? + .reshape((b, n, 3, self.num_heads, c / self.num_heads))? + .transpose(1, 2)? // 02134 + .transpose(0, 1)? // 20134 + .transpose(2, 3)?; // 20314 + let q = (qkv.i(0)? * self.scale)?; + let k = qkv.i(1)?.contiguous()?; + let v = qkv.i(2)?.contiguous()?; + let attn = (&q.matmul(&k.t()?)? + self._get_rel_pos_bias())?; + let attn = candle_nn::ops::softmax(&attn, D::Minus1)?; + let attn = attn.matmul(&v)?.transpose(1, 2)?.reshape((b, n, c))?; + self.proj.forward(&attn) + } +} + +#[derive(Debug)] +struct LayerScale { + gamma: Tensor, +} + +impl LayerScale { + fn new(vb: VarBuilder, dim: usize) -> Result { + let gamma = vb.get(dim, "gamma")?; + Ok(Self { gamma }) + } +} + +impl Module for LayerScale { + fn forward(&self, xs: &Tensor) -> Result { + xs.broadcast_mul(&self.gamma) + } +} + +#[derive(Debug)] +struct Mlp { + fc1: Linear, + fc2: Linear, +} + +impl Mlp { + fn new(vb: VarBuilder, in_features: usize, hidden_features: usize, bias: bool) -> Result { + let out_features = in_features; + let fc1 = linear(vb.pp("fc1"), in_features, hidden_features, bias)?; + let fc2 = linear(vb.pp("fc2"), hidden_features, out_features, bias)?; + Ok(Self { fc1, fc2 }) + } +} + +impl Module for Mlp { + fn forward(&self, xs: &Tensor) -> Result { + let xs = self.fc1.forward(xs)?.gelu()?; + self.fc2.forward(&xs) + } +} + +#[derive(Debug)] +struct Block { + norm1: LayerNorm, + attn: Attention, + ls1: LayerScale, + norm2: LayerNorm, + mlp: Mlp, + ls2: LayerScale, +} + +impl Block { + fn new( + vb: VarBuilder, + dim: usize, + num_heads: usize, + relative_position_index: &Tensor, + ) -> Result { + let norm1 = layer_norm(dim, 1e-6, vb.pp("norm1"))?; + let attn = Attention::new( + vb.pp("attn"), + dim, + num_heads, + true, + true, + relative_position_index, + )?; + let ls1 = LayerScale::new(vb.pp("ls1"), dim)?; + let norm2 = layer_norm(dim, 1e-6, vb.pp("norm2"))?; + let mlp = Mlp::new(vb.pp("mlp"), dim, dim * 4, true)?; + let ls2 = LayerScale::new(vb.pp("ls2"), dim)?; + Ok(Self { + norm1, + attn, + ls1, + norm2, + mlp, + ls2, + }) + } +} + +impl Module for Block { + fn forward(&self, xs: &Tensor) -> Result { + let residual = xs; + let xs = self + .ls1 + .forward(&self.attn.forward(&self.norm1.forward(xs)?)?)?; + let xs = (xs + residual)?; + let residual = &xs; + let xs = self + .ls2 + .forward(&self.mlp.forward(&self.norm2.forward(&xs)?)?)?; + xs + residual + } +} + +#[derive(Debug)] +struct PatchEmbed { + proj: candle_nn::Conv2d, + patch_size: (usize, usize), +} + +impl PatchEmbed { + fn new(vb: VarBuilder, patch_size: usize, in_chans: usize, embed_dim: usize) -> Result { + let config = candle_nn::Conv2dConfig { + stride: patch_size, + ..Default::default() + }; + let proj = candle_nn::conv2d(in_chans, embed_dim, patch_size, config, vb.pp("proj"))?; + Ok(Self { + proj, + patch_size: (patch_size, patch_size), + }) + } +} + +impl Module for PatchEmbed { + fn forward(&self, xs: &Tensor) -> Result { + let (_b, _c, h, w) = xs.dims4()?; + let (patch_h, patch_w) = self.patch_size; + if (h % patch_h) != 0 { + candle::bail!("image height {h} is not a multiple of patch height {patch_h}") + } + if (w % patch_w) != 0 { + candle::bail!("image width {w} is not a multiple of patch width {patch_w}") + } + let xs = self.proj.forward(xs)?; + let (b, c, h, w) = xs.dims4()?; + // flatten embeddings. + xs.reshape((b, c, h * w))?.transpose(1, 2) + } +} + +#[derive(Debug)] +pub struct BeitVisionTransformer { + patch_embed: PatchEmbed, + cls_token: Tensor, + blocks: Vec, + norm: LayerNorm, + head: Linear, +} + +impl BeitVisionTransformer { + pub fn new(vb: VarBuilder, depth: usize, embed_dim: usize, num_heads: usize) -> Result { + let patch_embed = PatchEmbed::new(vb.pp("patch_embed"), PATCH_SIZE, 3, embed_dim)?; + let cls_token = vb.get((1, 1, embed_dim), "cls_token")?; + let head = linear(vb.pp("head"), embed_dim, NUM_CLASSES, true)?; + let relative_position_index = vb.get((NB_TOKENS, NB_TOKENS), "relative_position_index")?; + let norm = layer_norm(embed_dim, 1e-6, vb.pp("norm"))?; + let vb_b = vb.pp("blocks"); + let blocks = (0..depth) + .map(|i| { + Block::new( + vb_b.pp(&i.to_string()), + embed_dim, + num_heads, + &relative_position_index, + ) + }) + .collect::>>()?; + Ok(Self { + patch_embed, + cls_token, + blocks, + norm, + head, + }) + } + + fn prepare_tokens_with_mask(&self, xs: &Tensor) -> Result { + let xs = self.patch_embed.forward(xs)?; + Tensor::cat(&[&self.cls_token, &xs], 1) + } + + fn get_intermediate_layers_not_chunked( + &self, + xs: &Tensor, + blocks_to_take: &[usize], + ) -> Result> { + let mut xs = self.prepare_tokens_with_mask(xs)?; + let mut output = Vec::new(); + for (i, blk) in self.blocks.iter().enumerate() { + xs = blk.forward(&xs)?; + if blocks_to_take.contains(&i) { + output.push(xs.clone()); + } + } + if output.len() != blocks_to_take.len() { + candle::bail!( + "only {} / {} blocks found", + output.len(), + blocks_to_take.len() + ); + } + Ok(output) + } + + pub fn get_intermediate_layers( + &self, + xs: &Tensor, + blocks_to_take: &[usize], + reshape: bool, + return_class_token: bool, + norm: bool, + ) -> Result { + let outputs = self.get_intermediate_layers_not_chunked(xs, blocks_to_take)?; + let outputs = if norm { + outputs + .iter() + .map(|out| self.norm.forward(out)) + .collect::>>()? + } else { + outputs + }; + let class_tokens = outputs + .iter() + .map(|out| out.i((.., 0))) + .collect::>>()?; + let outputs = outputs + .iter() + .map(|out| out.i((.., 1..))) + .collect::>>()?; + + let outputs = if reshape { + let (b, _c, w, h) = xs.dims4()?; + let patch_size = self.patch_embed.patch_size.0; + let num_channels = outputs[0].elem_count() / (b * (w / patch_size) * (h / patch_size)); + outputs + .iter() + .map(|out| { + out.reshape((b, w / patch_size, h / patch_size, num_channels))? + .transpose(2, 3)? + .transpose(1, 2) + }) + .collect::>>()? + } else { + outputs + }; + + let outputs = if return_class_token { + outputs + .iter() + .zip(class_tokens.iter()) + .map(|(out, class_token)| Tensor::cat(&[out, class_token], D::Minus1)) + .collect::>>()? + } else { + outputs + }; + + Tensor::stack(&outputs[..], 0) + } +} + +impl Module for BeitVisionTransformer { + fn forward(&self, xs: &Tensor) -> Result { + let mut xs = self.prepare_tokens_with_mask(xs)?; + for blk in self.blocks.iter() { + xs = blk.forward(&xs)? + } + let xs_moy_local_tokens = xs.i((.., 1..))?.mean(1)?; + let xs_norm = self.norm.forward(&xs_moy_local_tokens)?; + self.head.forward(&xs_norm) + } +} + +pub fn vit_base(vb: VarBuilder) -> Result { + BeitVisionTransformer::new(vb, 12, 768, 12) +} + +pub fn vit_large(vb: VarBuilder) -> Result { + BeitVisionTransformer::new(vb, 24, 1024, 16) +} diff --git a/candle-transformers/src/models/mod.rs b/candle-transformers/src/models/mod.rs index 2908d3457a..d95d30ae5a 100644 --- a/candle-transformers/src/models/mod.rs +++ b/candle-transformers/src/models/mod.rs @@ -1,3 +1,4 @@ +pub mod beit; pub mod bert; pub mod bigcode; pub mod blip;