How can I implement Max Pooling in Arrayfire in rust without resorting to writing my own cuda code

Question

I'm trying to figure out how to implement max pooling on Arrayfire. My current best approach involves iterating over each convolved output and apply a function which applies four kernels, [1 0 0 0], [0 1 0 0], [0 0 1 0], [0 0 0 1], and produces four outputs, which I can then compare for the maximum value at each pixel.

My issue with that is it seems terribly slow and incorrect to be looping like that in a tensor library, but I haven't been able to come up with a better solution

Answer 1

I have settled on the following: index the quandrants with seq, then grab the maximums

#[test]
    fn maxfilt____() {
        let fourxfour = Array::new(&(0..16).into_iter().collect::<Vec<_>>(), dim4!(4, 4, 1, 1));

        let dim0 = fourxfour.dims()[0] as i32;
        let dim1 = fourxfour.dims()[1] as i32;

        let q1_indices = &[seq!(0, dim0 - 1, 2), seq!(0, dim1 - 1, 2), seq!(), seq!()];
        let q2_indices = &[seq!(0, dim0 - 1, 2), seq!(1, dim1 - 1, 2), seq!(), seq!()];
        let q3_indices = &[seq!(1, dim0 - 1, 2), seq!(0, dim1 - 1, 2), seq!(), seq!()];
        let q4_indices = &[seq!(1, dim0 - 1, 2), seq!(1, dim1 - 1, 2), seq!(), seq!()];

        let q1s = index(&fourxfour, q1_indices);
        let q2s = index(&fourxfour, q2_indices);
        let q3s = index(&fourxfour, q3_indices);
        let q4s = index(&fourxfour, q4_indices);
   
        let max = maxof(&q1s, &maxof(&q2s, &maxof(&q3s, &q4s, false), false), false);

        af_print!("max", max);
    }