candle_transformers/models/
beit.rs

1//! Based on the BEIT vision-language model.
2//!
3//! See "BEIT: BERT Pre-Training of Image Transformers", Bao et al. 2021
4//! - [Arxiv](https://arxiv.org/abs/2106.08254)
5//! - [Github](https://github.com/microsoft/unilm/tree/master/beit)
6//!
7
8use candle::{DType, Device, IndexOp, Result, Tensor, D};
9use candle_nn::{layer_norm, LayerNorm, Linear, Module, VarBuilder};
10
11const IMG_SIZE: usize = 384;
12const PATCH_SIZE: usize = 16;
13const NUM_CLASSES: usize = 1000;
14const WINDOW_SIZE: usize = IMG_SIZE / PATCH_SIZE; // 384 / 16 = 24
15const NB_TOKENS: usize = WINDOW_SIZE * WINDOW_SIZE + 1; // 24 * 24 + 1 = 577
16
17fn linear(vb: VarBuilder, in_dim: usize, out_dim: usize, bias: bool) -> Result<Linear> {
18    if bias {
19        candle_nn::linear(in_dim, out_dim, vb)
20    } else {
21        candle_nn::linear_no_bias(in_dim, out_dim, vb)
22    }
23}
24
25#[derive(Debug)]
26struct Attention {
27    qkv: Linear,
28    proj: Linear,
29    relative_position_bias_table: Tensor,
30    relative_position_index: Tensor,
31    num_heads: usize,
32    scale: f64,
33}
34
35impl Attention {
36    fn new(
37        vb: VarBuilder,
38        dim: usize,
39        num_heads: usize,
40        qkv_bias: bool,
41        proj_bias: bool,
42    ) -> Result<Self> {
43        let qkv = linear(vb.pp("qkv"), dim, dim * 3, qkv_bias)?;
44        let proj = linear(vb.pp("proj"), dim, dim, proj_bias)?;
45        // num_relative_distance = token-token(47x47) + token-CLS(1) + CLS-token(1) + CLS-CLS(1) = 2212
46        let num_relative_distance = (2 * WINDOW_SIZE - 1) * (2 * WINDOW_SIZE - 1) + 3;
47        let relative_position_bias_table = vb.get(
48            (num_relative_distance, num_heads),
49            "relative_position_bias_table",
50        )?;
51        let relative_position_index =
52            Self::gen_relative_position_index(relative_position_bias_table.device())?;
53        let scale = 1. / ((dim / num_heads) as f64).sqrt();
54        Ok(Self {
55            qkv,
56            proj,
57            relative_position_bias_table,
58            relative_position_index,
59            num_heads,
60            scale,
61        })
62    }
63}
64
65impl Attention {
66    // See: https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/beit.py#L61
67    fn gen_relative_position_index(device: &Device) -> Result<Tensor> {
68        let num_relative_distance = (2 * WINDOW_SIZE - 1) * (2 * WINDOW_SIZE - 1) + 3;
69        let w_area = WINDOW_SIZE * WINDOW_SIZE;
70
71        let t_arange: Tensor = Tensor::arange(0, WINDOW_SIZE as u32, device)?;
72        let t_ndgrid = Tensor::meshgrid(&[&t_arange, &t_arange], false)?;
73        let coords_flatten = Tensor::stack(&t_ndgrid, 0)?.flatten(1, 2)?;
74
75        let tmp1 = coords_flatten
76            .unsqueeze(2)?
77            .broadcast_as((2, w_area, w_area))?
78            .to_dtype(DType::I64)?;
79        let tmp2 = coords_flatten
80            .unsqueeze(1)?
81            .broadcast_as((2, w_area, w_area))?
82            .to_dtype(DType::I64)?;
83        let relative_coords = (tmp1 - tmp2)?
84            .transpose(0, 1)? // 102
85            .transpose(1, 2)? // 120
86            .contiguous()?;
87
88        let relative_coords = relative_coords.slice_assign(
89            &[0..w_area, 0..w_area, 0..1],
90            &(relative_coords.i((0..w_area, 0..w_area, 0..1))? + (WINDOW_SIZE - 1) as f64)?,
91        )?;
92        let relative_coords = relative_coords.slice_assign(
93            &[0..w_area, 0..w_area, 1..2],
94            &(relative_coords.i((0..w_area, 0..w_area, 1..2))? + (WINDOW_SIZE - 1) as f64)?,
95        )?;
96        let relative_coords = relative_coords.slice_assign(
97            &[0..w_area, 0..w_area, 0..1],
98            &(relative_coords.i((.., .., 0..1))? * (2. * (WINDOW_SIZE as f64) - 1.))?,
99        )?;
100
101        Tensor::zeros((w_area + 1, w_area + 1), DType::I64, device)?
102            .slice_assign(&[1.., 1..], &relative_coords.sum(2)?)?
103            .slice_assign(
104                &[0..1, 0..(w_area + 1)],
105                &(Tensor::ones((1, w_area + 1), DType::I64, device)?
106                    * ((num_relative_distance - 3) as f64))?
107                    .to_dtype(DType::I64)?,
108            )?
109            .slice_assign(
110                &[0..(w_area + 1), 0..1],
111                &(Tensor::ones((w_area + 1, 1), DType::I64, device)?
112                    * ((num_relative_distance - 2) as f64))?
113                    .to_dtype(DType::I64)?,
114            )?
115            .slice_assign(
116                &[0..1, 0..1],
117                &(Tensor::ones((1, 1), DType::I64, device)?
118                    * ((num_relative_distance - 1) as f64))?
119                    .to_dtype(DType::I64)?,
120            )
121    }
122
123    fn _get_rel_pos_bias(&self) -> Result<Tensor> {
124        self.relative_position_bias_table
125            .index_select(
126                &self
127                    .relative_position_index
128                    .flatten_all()?
129                    .to_dtype(DType::U32)?,
130                0,
131            )?
132            .reshape((NB_TOKENS, NB_TOKENS, ()))?
133            .transpose(0, 1)? // 102
134            .transpose(0, 2)? // 201
135            .contiguous()?
136            .unsqueeze(0)
137    }
138}
139
140impl Module for Attention {
141    fn forward(&self, xs: &Tensor) -> Result<Tensor> {
142        let (b, n, c) = xs.dims3()?;
143        let qkv = self
144            .qkv
145            .forward(xs)?
146            .reshape((b, n, 3, self.num_heads, c / self.num_heads))?
147            .transpose(1, 2)? // 02134
148            .transpose(0, 1)? // 20134
149            .transpose(2, 3)?; // 20314
150        let q = (qkv.i(0)? * self.scale)?;
151        let k = qkv.i(1)?.contiguous()?;
152        let v = qkv.i(2)?.contiguous()?;
153        let attn = (&q.matmul(&k.t()?)? + self._get_rel_pos_bias())?;
154        let attn = candle_nn::ops::softmax(&attn, D::Minus1)?;
155        let attn = attn.matmul(&v)?.transpose(1, 2)?.reshape((b, n, c))?;
156        self.proj.forward(&attn)
157    }
158}
159
160#[derive(Debug)]
161struct LayerScale {
162    gamma: Tensor,
163}
164
165impl LayerScale {
166    fn new(vb: VarBuilder, dim: usize) -> Result<Self> {
167        let gamma = vb.get(dim, "gamma")?;
168        Ok(Self { gamma })
169    }
170}
171
172impl Module for LayerScale {
173    fn forward(&self, xs: &Tensor) -> Result<Tensor> {
174        xs.broadcast_mul(&self.gamma)
175    }
176}
177
178#[derive(Debug)]
179struct Mlp {
180    fc1: Linear,
181    fc2: Linear,
182}
183
184impl Mlp {
185    fn new(vb: VarBuilder, in_features: usize, hidden_features: usize, bias: bool) -> Result<Self> {
186        let out_features = in_features;
187        let fc1 = linear(vb.pp("fc1"), in_features, hidden_features, bias)?;
188        let fc2 = linear(vb.pp("fc2"), hidden_features, out_features, bias)?;
189        Ok(Self { fc1, fc2 })
190    }
191}
192
193impl Module for Mlp {
194    fn forward(&self, xs: &Tensor) -> Result<Tensor> {
195        let xs = self.fc1.forward(xs)?.gelu()?;
196        self.fc2.forward(&xs)
197    }
198}
199
200#[derive(Debug)]
201struct Block {
202    norm1: LayerNorm,
203    attn: Attention,
204    ls1: LayerScale,
205    norm2: LayerNorm,
206    mlp: Mlp,
207    ls2: LayerScale,
208}
209
210impl Block {
211    fn new(vb: VarBuilder, dim: usize, num_heads: usize) -> Result<Self> {
212        let norm1 = layer_norm(dim, 1e-6, vb.pp("norm1"))?;
213        let attn = Attention::new(vb.pp("attn"), dim, num_heads, true, true)?;
214        let ls1 = LayerScale::new(vb.pp("ls1"), dim)?;
215        let norm2 = layer_norm(dim, 1e-6, vb.pp("norm2"))?;
216        let mlp = Mlp::new(vb.pp("mlp"), dim, dim * 4, true)?;
217        let ls2 = LayerScale::new(vb.pp("ls2"), dim)?;
218        Ok(Self {
219            norm1,
220            attn,
221            ls1,
222            norm2,
223            mlp,
224            ls2,
225        })
226    }
227}
228
229impl Module for Block {
230    fn forward(&self, xs: &Tensor) -> Result<Tensor> {
231        let residual = xs;
232        let xs = self
233            .ls1
234            .forward(&self.attn.forward(&self.norm1.forward(xs)?)?)?;
235        let xs = (xs + residual)?;
236        let residual = &xs;
237        let xs = self
238            .ls2
239            .forward(&self.mlp.forward(&self.norm2.forward(&xs)?)?)?;
240        xs + residual
241    }
242}
243
244#[derive(Debug)]
245struct PatchEmbed {
246    proj: candle_nn::Conv2d,
247    patch_size: (usize, usize),
248}
249
250impl PatchEmbed {
251    fn new(vb: VarBuilder, patch_size: usize, in_chans: usize, embed_dim: usize) -> Result<Self> {
252        let config = candle_nn::Conv2dConfig {
253            stride: patch_size,
254            ..Default::default()
255        };
256        let proj = candle_nn::conv2d(in_chans, embed_dim, patch_size, config, vb.pp("proj"))?;
257        Ok(Self {
258            proj,
259            patch_size: (patch_size, patch_size),
260        })
261    }
262}
263
264impl Module for PatchEmbed {
265    fn forward(&self, xs: &Tensor) -> Result<Tensor> {
266        let (_b, _c, h, w) = xs.dims4()?;
267        let (patch_h, patch_w) = self.patch_size;
268        if (h % patch_h) != 0 {
269            candle::bail!("image height {h} is not a multiple of patch height {patch_h}")
270        }
271        if (w % patch_w) != 0 {
272            candle::bail!("image width {w} is not a multiple of patch width {patch_w}")
273        }
274        let xs = self.proj.forward(xs)?;
275        let (b, c, h, w) = xs.dims4()?;
276        // flatten embeddings.
277        xs.reshape((b, c, h * w))?.transpose(1, 2)
278    }
279}
280
281#[derive(Debug)]
282pub struct BeitVisionTransformer {
283    patch_embed: PatchEmbed,
284    cls_token: Tensor,
285    blocks: Vec<Block>,
286    norm: LayerNorm,
287    head: Linear,
288}
289
290impl BeitVisionTransformer {
291    pub fn new(vb: VarBuilder, depth: usize, embed_dim: usize, num_heads: usize) -> Result<Self> {
292        let patch_embed = PatchEmbed::new(vb.pp("patch_embed"), PATCH_SIZE, 3, embed_dim)?;
293        let cls_token = vb.get((1, 1, embed_dim), "cls_token")?;
294        let head = linear(vb.pp("head"), embed_dim, NUM_CLASSES, true)?;
295        let norm = layer_norm(embed_dim, 1e-6, vb.pp("norm"))?;
296        let vb_b = vb.pp("blocks");
297        let blocks = (0..depth)
298            .map(|i| Block::new(vb_b.pp(i.to_string()), embed_dim, num_heads))
299            .collect::<Result<Vec<_>>>()?;
300        Ok(Self {
301            patch_embed,
302            cls_token,
303            blocks,
304            norm,
305            head,
306        })
307    }
308
309    fn prepare_tokens_with_mask(&self, xs: &Tensor) -> Result<Tensor> {
310        let xs = self.patch_embed.forward(xs)?;
311        Tensor::cat(&[&self.cls_token, &xs], 1)
312    }
313
314    fn get_intermediate_layers_not_chunked(
315        &self,
316        xs: &Tensor,
317        blocks_to_take: &[usize],
318    ) -> Result<Vec<Tensor>> {
319        let mut xs = self.prepare_tokens_with_mask(xs)?;
320        let mut output = Vec::new();
321        for (i, blk) in self.blocks.iter().enumerate() {
322            xs = blk.forward(&xs)?;
323            if blocks_to_take.contains(&i) {
324                output.push(xs.clone());
325            }
326        }
327        if output.len() != blocks_to_take.len() {
328            candle::bail!(
329                "only {} / {} blocks found",
330                output.len(),
331                blocks_to_take.len()
332            );
333        }
334        Ok(output)
335    }
336
337    pub fn get_intermediate_layers(
338        &self,
339        xs: &Tensor,
340        blocks_to_take: &[usize],
341        reshape: bool,
342        return_class_token: bool,
343        norm: bool,
344    ) -> Result<Tensor> {
345        let outputs = self.get_intermediate_layers_not_chunked(xs, blocks_to_take)?;
346        let outputs = if norm {
347            outputs
348                .iter()
349                .map(|out| self.norm.forward(out))
350                .collect::<Result<Vec<_>>>()?
351        } else {
352            outputs
353        };
354        let class_tokens = outputs
355            .iter()
356            .map(|out| out.i((.., 0)))
357            .collect::<Result<Vec<_>>>()?;
358        let outputs = outputs
359            .iter()
360            .map(|out| out.i((.., 1..)))
361            .collect::<Result<Vec<_>>>()?;
362
363        let outputs = if reshape {
364            let (b, _c, w, h) = xs.dims4()?;
365            let patch_size = self.patch_embed.patch_size.0;
366            let num_channels = outputs[0].elem_count() / (b * (w / patch_size) * (h / patch_size));
367            outputs
368                .iter()
369                .map(|out| {
370                    out.reshape((b, w / patch_size, h / patch_size, num_channels))?
371                        .transpose(2, 3)?
372                        .transpose(1, 2)
373                })
374                .collect::<Result<Vec<_>>>()?
375        } else {
376            outputs
377        };
378
379        let outputs = if return_class_token {
380            outputs
381                .iter()
382                .zip(class_tokens.iter())
383                .map(|(out, class_token)| Tensor::cat(&[out, class_token], D::Minus1))
384                .collect::<Result<Vec<_>>>()?
385        } else {
386            outputs
387        };
388
389        Tensor::stack(&outputs[..], 0)
390    }
391}
392
393impl Module for BeitVisionTransformer {
394    fn forward(&self, xs: &Tensor) -> Result<Tensor> {
395        let mut xs = self.prepare_tokens_with_mask(xs)?;
396        for blk in self.blocks.iter() {
397            xs = blk.forward(&xs)?
398        }
399        let xs_moy_local_tokens = xs.i((.., 1..))?.mean(1)?;
400        let xs_norm = self.norm.forward(&xs_moy_local_tokens)?;
401        self.head.forward(&xs_norm)
402    }
403}
404
405pub fn vit_base(vb: VarBuilder) -> Result<BeitVisionTransformer> {
406    BeitVisionTransformer::new(vb, 12, 768, 12)
407}
408
409pub fn vit_large(vb: VarBuilder) -> Result<BeitVisionTransformer> {
410    BeitVisionTransformer::new(vb, 24, 1024, 16)
411}
candle_transformers/models/beit.rs

candle_transformers/models/
beit.rs
