style: cleanup resampler-pos

Author: CearX

models/clip/common-cpu/src/infer.rs

@@ -96,3 +96,163 @@ fn pos70(n: usize, [w, h]: [usize; 2], d_patch: usize) -> Tensor<Blob> {
 
     ans
 }
+
+#[cfg(test)]
+mod test_pos_embd {
+    use super::*;
+
+    mod c_like {
+        use super::*;
+        pub(super) fn pos_resampler(n: usize, [w, h]: [usize; 2], d_patch: usize) -> Tensor<Blob> {
+            let d = 3584;
+            let pos_w = w / d_patch;
+            let pos_h = h / d_patch;
+
+            let mut ans = Tensor::new(ty::F32, &[1, pos_w * pos_h, d])
+                .broadcast(0, n)
+                .map(Blob::new);
+            let (&mut [], data, &mut []) = (unsafe { ans.get_mut().align_to_mut::<f32>() }) else {
+                panic!()
+            };
+
+            let pos_embed_t = get_2d_sincos_pos_embed(d, (pos_w, pos_h));
+
+            for i in 0..pos_w * pos_h {
+                for j in 0..d {
+                    data[i * d + j] = pos_embed_t[i][j];
+                }
+            }
+            ans
+        }
+
+        fn get_2d_sincos_pos_embed(embed_dim: usize, image_size: (usize, usize)) -> Vec<Vec<f32>> {
+            let (grid_h_size, grid_w_size) = image_size;
+
+            let mut grid_h: Vec<f32> = (0..grid_h_size).map(|i| i as f32).collect();
+            let mut grid_w: Vec<f32> = (0..grid_w_size).map(|i| i as f32).collect();
+
+            let mut grid: Vec<Vec<f32>> = vec![vec![0.0; grid_w_size]; grid_h_size];
+            for h in 0..grid_h_size {
+                for w in 0..grid_w_size {
+                    grid[h][w] = grid_w[w];
+                }
+            }
+
+            let mut grid_2d: Vec<Vec<Vec<f32>>> = vec![grid.clone(), grid.clone()];
+            for h in 0..grid_h_size {
+                for w in 0..grid_w_size {
+                    grid_2d[0][h][w] = grid_h[h];
+                    grid_2d[1][h][w] = grid_w[w];
+                }
+            }
+
+            let pos_embed_3d = get_2d_sincos_pos_embed_from_grid(embed_dim, grid_2d);
+
+            let (H, W) = image_size;
+            let mut pos_embed_2d: Vec<Vec<f32>> = vec![vec![0.0; embed_dim]; H * W];
+            for h in 0..H {
+                for w in 0..W {
+                    pos_embed_2d[w * H + h] = pos_embed_3d[h][w].clone();
+                }
+            }
+
+            pos_embed_2d
+        }
+
+        fn get_2d_sincos_pos_embed_from_grid(
+            embed_dim: usize,
+            grid: Vec<Vec<Vec<f32>>>,
+        ) -> Vec<Vec<Vec<f32>>> {
+            assert!(embed_dim % 2 == 0);
+
+            let emb_h = get_1d_sincos_pos_embed_from_grid_new(embed_dim / 2, grid[0].clone()); // (H, W, D/2)
+            let emb_w = get_1d_sincos_pos_embed_from_grid_new(embed_dim / 2, grid[1].clone()); // (H, W, D/2)
+
+            let H = emb_h.len();
+            let W = emb_h[0].len();
+            let mut emb: Vec<Vec<Vec<f32>>> = vec![vec![vec![0.0; embed_dim]; W]; H];
+
+            for h in 0..H {
+                for w in 0..W {
+                    for d in 0..(embed_dim / 2) {
+                        emb[h][w][d] = emb_h[h][w][d];
+                        emb[h][w][d + embed_dim / 2] = emb_w[h][w][d];
+                    }
+                }
+            }
+
+            emb
+        }
+
+        fn get_1d_sincos_pos_embed_from_grid_new(
+            embed_dim: usize,
+            pos: Vec<Vec<f32>>,
+        ) -> Vec<Vec<Vec<f32>>> {
+            assert!(embed_dim % 2 == 0);
+            let H = pos.len();
+            let W = pos[0].len();
+
+            let mut omega: Vec<f32> = (0..embed_dim / 2)
+                .map(|i| 1.0 / 10000.0f32.powi(i as i32 / (embed_dim / 2) as i32))
+                .collect();
+
+            let mut emb: Vec<Vec<Vec<f32>>> = vec![vec![vec![0.0; embed_dim]; W]; H];
+            for h in 0..H {
+                for w in 0..W {
+                    for d in 0..(embed_dim / 2) {
+                        let out_value = pos[h][w] * omega[d];
+                        emb[h][w][d] = out_value.sin();
+                        emb[h][w][d + embed_dim / 2] = out_value.cos();
+                    }
+                }
+            }
+
+            emb
+        }
+    }
+
+    mod rust_style {
+        use super::*;
+        pub(super) fn pos_resampler(n: usize, [w, h]: [usize; 2], d_patch: usize) -> Tensor<Blob> {
+            let d = 3584;
+            assert!(d % 4 == 0);
+
+            let pos_w = w / d_patch;
+            let pos_h = h / d_patch;
+
+            let mut ans = Tensor::new(ty::F32, &[1, pos_w * pos_h, d]).map(Blob::new);
+            let (&mut [], data, &mut []) = (unsafe { ans.get_mut().align_to_mut::<f32>() }) else {
+                panic!()
+            };
+            set_2d_sincos_pos_embed(data, d, (pos_w, pos_h));
+
+            ans.broadcast(0, n)
+        }
+
+        fn set_2d_sincos_pos_embed(data: &mut [f32], d: usize, (h, w): (usize, usize)) {
+            for r in 0..h {
+                for c in 0..w {
+                    let data = &mut data[(c * h + r) * d..][..d];
+                    let d = d / 4;
+
+                    for i in 0..d {
+                        let (sin, cos) = (r as f32).sin_cos();
+                        data[0 * d..][i] = sin;
+                        data[1 * d..][i] = cos;
+
+                        let (sin, cos) = (c as f32).sin_cos();
+                        data[2 * d..][i] = sin;
+                        data[3 * d..][i] = cos;
+                    }
+                }
+            }
+        }
+    }
+
+    #[test]
+    fn test_eq() {
+        let a = c_like::pos_resampler(4, [336, 224], 14).take();
+        let b = rust_style::pos_resampler(4, [336, 224], 14).take();
+        assert_eq!(&*a, &*b);
+    }
+}