Conv2d doesn't work on rectangular image about burn

[antimora]

It would be nice if the Tensor documentation had this explanation.

The dimensions of forward pass for Conv2d are actually described in the API docs:
https://burn.dev/docs/burn/nn/conv/struct.Conv2d.html#method.forward

[nathanielsimard]

I'm not sure to understand the issue, there are a lot of configurations possible with conv2d and they won't return the same result, it all depends on your use case. If you think there is a bug, provide a scenario that we can reproduce with numbers (images won't help much 😅).

[J-F-Liu]

Please try the following code with different images, the expected output is blured image, rather than disordered.

use burn::{
    backend::ndarray::{NdArray, NdArrayDevice},
    module::{ConstantRecord, Param, ParamId},
    nn::{self, conv::Conv2dRecord, PaddingConfig2d},
    tensor::{backend::Backend, Tensor},
};
use image::ImageEncoder;

fn main() {
    type Backend = NdArray;
    let device = NdArrayDevice::Cpu;
    let (input, (width, height)) = read_image("umbralla.png").expect("read image");
    let input = input.into_iter().map(|p| p as f32).collect::<Vec<_>>();
    let input: Tensor<Backend, 4> = Tensor::from_floats(input.as_slice(), &device).reshape([
        1,
        1,
        width as usize,
        height as usize,
    ]);
    let (kernel, size) = gaussian_kernel(0.65);
    let guassian = create_conv2d(
        [1, 1],
        [size, size],
        &kernel,
        PaddingConfig2d::Same,
        &device,
    );
    dbg!(&guassian);
    let output = guassian.forward(input);
    let output = output
        .reshape([width as usize, height as usize])
        .into_data();
    save_image(&output, "smoothed.png").expect("save image");
}

pub fn read_image<P: AsRef<std::path::Path>>(
    filepath: P,
) -> image::ImageResult<(Vec<u8>, (u32, u32))> {
    let image = image::open(filepath)?;
    let gray_image = match image {
        image::DynamicImage::ImageLuma8(gray_image) => gray_image,
        _ => image.into_luma8(),
    };
    let size = (gray_image.width(), gray_image.height());
    let content = gray_image.into_raw();
    Ok((content, size))
}

pub fn save_image<P: AsRef<std::path::Path>>(
    data: &burn::tensor::Data<f32, 2>,
    filepath: P,
) -> image::ImageResult<()> {
    let [width, height] = data.shape.dims;
    let pixels = data
        .value
        .iter()
        .map(|p| p.round() as u8)
        .collect::<Vec<_>>();
    let file = std::fs::File::create(filepath)?;
    let buffer = std::io::BufWriter::new(file);
    let encoder = image::codecs::png::PngEncoder::new(buffer);
    encoder.write_image(&pixels, width as u32, height as u32, image::ColorType::L8)
}

fn create_conv2d<B: Backend>(
    channels: [usize; 2],
    kernel_size: [usize; 2],
    kernel: &[f32],
    padding: PaddingConfig2d,
    device: &B::Device,
) -> nn::conv::Conv2d<B> {
    let record = Conv2dRecord {
        weight: Param::new(
            ParamId::new(),
            Tensor::from_floats(kernel, device).reshape([
                channels[1], // channels_out
                channels[0], // channels_in / groups
                kernel_size[0],
                kernel_size[1],
            ]),
        ),
        bias: None,
        stride: [ConstantRecord::new(), ConstantRecord::new()],
        kernel_size: [ConstantRecord::new(), ConstantRecord::new()],
        dilation: [ConstantRecord::new(), ConstantRecord::new()],
        groups: ConstantRecord::new(),
        padding: ConstantRecord::new(),
    };
    nn::conv::Conv2dConfig::new(channels, kernel_size)
        .with_bias(false)
        .with_groups(1)
        .with_padding(padding)
        .with_stride([1, 1])
        .init_with(record)
}

/// compute gaussian kernel e^(-(x^2+y^2)/(2σ^2))
/// return (kernel, size), kernel size is `2 * radius + 1`
fn gaussian_kernel(sigma: f32) -> (Vec<f32>, usize) {
    /*
      The size of the kernel is selected to guarantee that the first discarded
      term is at least 10^prec times smaller than the central value. For that,
      the half size of the kernel must be larger than x, with
        e^(-x^2/2sigma^2) = 1/10^prec
      Then,
        x = sigma * sqrt( 2 * prec * ln(10) )
    */
    let prec = 3.0;
    let radius = (sigma * (2.0 * prec * 10.0_f32.ln()).sqrt()).ceil() as i32;
    let size = 2 * radius + 1; /* kernel size */
    let mut kernel = Vec::with_capacity((size * size) as usize);
    for y in 0..size {
        for x in 0..size {
            let dist2 = (x - radius).pow(2) + (y - radius).pow(2);
            // proximate a circular region
            let value = if dist2 <= radius * radius {
                (-0.5 * (dist2 as f32) / sigma.powi(2)).exp()
            } else {
                0.0
            };
            kernel.push(value);
        }
    }

    //normalization
    let sum: f32 = kernel.iter().sum();
    if sum >= 0.0 {
        for elem in kernel.iter_mut() {
            *elem /= sum;
        }
    }

    (kernel, size as usize)
}

[TomWyllie]

I think your problem might be to do with the order of your dimensions. The diagonal artefacts in your image and the fact it works when the image is square (i.e. [height, width] == [width, height] ) suggest you might be interpreting the order of the 2D data stored in a 1D buffer incorrectly.

In the documentation for Conv2D, it does say

• input: [batch_size, channels_in, height_in, width_in],
• output: [batch_size, channels_out, height_out, width_out],

You have width first then height, so perhaps this is causing your problem.

[J-F-Liu]

@TomWyllie Thanks, after swap width and height the program works, so is the Tensor treat image rows as matrix columns, elements are assumed in column major order?

[TomWyllie]

The elements are actually still in row major order - if the shape is [batch, channels, height, width], then the last dimension is the width dimension. Since elements are stored contiguously in memory along last dimension to first, this layout is row major. I got quite confused by this at first, but what has helped me is realising that the [height, width] convention is the same as [num_rows, num_cols] which is the same "rows then columns" convention as regular matrices. Again, this is row major because storing the number of columns in the last dimension means neighbouring elements are in neighbouring columns (in the same row). I put together a short snippet in case it's useful to show the row major tensor and how this works with the op, which has the same convention:

    let x_f: Vec<f32> = vec![1., 2., 3., 4., 5., 6.];
    let x: Tensor<B, 1> = Tensor::from_floats(x_f.as_slice(), &B::Device::default());

    //  batch, channels, height, width
    let x = x.reshape(Shape::new([1, 1, 2, 3]));
    println!("{:?}", x);
    //  [[[[ 1. , 2. , 3. ],
    //     [ 4. , 5. , 6. ]]]]

    // 1 channel -> 1 channel, kernel size 2x1 = 2 rows x 1 col = height 2, width 1
    let conv: Conv2d<B> = Conv2dConfig::new([1, 1], [2, 1])
        .with_initializer(Initializer::Ones)
        .with_bias(false)
        .init(&B::Device::default());

    let y = conv.forward(x);
    println!("{:?}", y);
    //  [[[[ 5. , 7. , 9. ]]]]

[J-F-Liu]

It would be nice if the Tensor documentation had this explanation.