burn_nn/modules/conv/
deform_conv2d.rs

1use alloc::format;
2use burn::tensor::ops::DeformConvOptions;
3
4use burn_core as burn;
5
6use crate::PaddingConfig2d;
7use burn::config::Config;
8use burn::module::Initializer;
9use burn::module::{Content, DisplaySettings, Module, ModuleDisplay, Param};
10use burn::tensor::Tensor;
11use burn::tensor::backend::Backend;
12use burn::tensor::module::deform_conv2d;
13
14use crate::conv::checks;
15
16/// Configuration to create a [deformable 2D convolution](DeformConv2d) layer, using the [init function](DeformConv2dConfig::init).
17#[derive(Config, Debug)]
18pub struct DeformConv2dConfig {
19    /// The number of channels.
20    pub channels: [usize; 2],
21    /// The size of the kernel.
22    pub kernel_size: [usize; 2],
23    /// The stride of the convolution.
24    #[config(default = "[1, 1]")]
25    pub stride: [usize; 2],
26    /// Spacing between kernel elements.
27    #[config(default = "[1, 1]")]
28    pub dilation: [usize; 2],
29    /// Controls the connections between input and output channels.
30    #[config(default = "1")]
31    pub weight_groups: usize,
32    /// Offset groups.
33    #[config(default = "1")]
34    pub offset_groups: usize,
35    /// The padding configuration.
36    ///
37    /// ### Warning
38    /// Only symmetric padding is currently supported. As such, using `Same` padding with an even kernel
39    /// size is not supported as it will not produce the same output size.
40    #[config(default = "PaddingConfig2d::Valid")]
41    pub padding: PaddingConfig2d,
42    /// If bias should be added to the output.
43    #[config(default = true)]
44    pub bias: bool,
45    /// The type of function used to initialize neural network parameters
46    #[config(
47        default = "Initializer::KaimingUniform{gain:1.0/num_traits::Float::sqrt(3.0),fan_out_only:false}"
48    )]
49    pub initializer: Initializer,
50}
51
52/// Applies a deformable 2D convolution over input tensors.
53///
54/// Should be created with [DeformConv2dConfig].
55#[derive(Module, Debug)]
56#[module(custom_display)]
57pub struct DeformConv2d<B: Backend> {
58    /// Tensor of shape `[channels_out, channels_in / groups, kernel_size_1, kernel_size_2]`
59    pub weight: Param<Tensor<B, 4>>,
60    /// Tensor of shape `[channels_out]`
61    pub bias: Option<Param<Tensor<B, 1>>>,
62    /// Stride of the convolution.
63    pub stride: [usize; 2],
64    /// Size of the kernel.
65    pub kernel_size: [usize; 2],
66    /// Spacing between kernel elements.
67    pub dilation: [usize; 2],
68    /// Controls the connections between input and output channels.
69    pub weight_groups: usize,
70    /// Offset groups.
71    pub offset_groups: usize,
72    /// The padding configuration.
73    pub padding: PaddingConfig2d,
74}
75
76impl DeformConv2dConfig {
77    /// Initialize a new [DeformConv2d](DeformConv2d) module.
78    pub fn init<B: Backend>(&self, device: &B::Device) -> DeformConv2d<B> {
79        checks::checks_channels_div_groups(self.channels[0], self.channels[1], self.weight_groups);
80        if self.padding == PaddingConfig2d::Same {
81            checks::check_same_padding_support(&self.kernel_size);
82        }
83
84        let shape = [
85            self.channels[1],
86            self.channels[0] / self.weight_groups,
87            self.kernel_size[0],
88            self.kernel_size[1],
89        ];
90
91        let k = self.kernel_size.iter().product::<usize>();
92        let fan_in = self.channels[0] / self.weight_groups * k;
93        let fan_out = self.channels[1] / self.weight_groups * k;
94
95        let weight = self
96            .initializer
97            .init_with(shape, Some(fan_in), Some(fan_out), device);
98        let mut bias = None;
99
100        if self.bias {
101            bias = Some(self.initializer.init_with(
102                [self.channels[1]],
103                Some(fan_in),
104                Some(fan_out),
105                device,
106            ));
107        }
108
109        DeformConv2d {
110            weight,
111            bias,
112            stride: self.stride,
113            kernel_size: self.kernel_size,
114            dilation: self.dilation,
115            padding: self.padding.clone(),
116            weight_groups: self.weight_groups,
117            offset_groups: self.weight_groups,
118        }
119    }
120}
121
122impl<B: Backend> ModuleDisplay for DeformConv2d<B> {
123    fn custom_settings(&self) -> Option<DisplaySettings> {
124        DisplaySettings::new()
125            .with_new_line_after_attribute(false)
126            .optional()
127    }
128
129    fn custom_content(&self, content: Content) -> Option<Content> {
130        // Format the stride, kernel_size and dilation as strings, formatted as arrays instead of indexed.
131        let stride = format!("{:?}", self.stride);
132        let kernel_size = format!("{:?}", self.kernel_size);
133        let dilation = format!("{:?}", self.dilation);
134
135        content
136            .add("stride", &stride)
137            .add("kernel_size", &kernel_size)
138            .add("dilation", &dilation)
139            .add("weight_groups", &self.weight_groups)
140            .add("offset_groups", &self.offset_groups)
141            .add_debug_attribute("padding", &self.padding)
142            .optional()
143    }
144}
145
146impl<B: Backend> DeformConv2d<B> {
147    /// Applies the forward pass on the input tensor.
148    ///
149    /// See [deform_conv2d](burn::tensor::module::deform_conv2d) for more information.
150    ///
151    /// # Shapes
152    ///
153    /// - input: `[batch_size, channels_in, height_in, width_in]`
154    /// - offset: `[batch_size, 2 * offset_groups * kernel_height * kernel_width, height_out, width_out]`
155    /// - mask: `[batch_size, offset_groups * kernel_height * kernel_width, height_out, width_out]`
156    /// - output: `[batch_size, channels_out, height_out, width_out]`
157    pub fn forward(
158        &self,
159        input: Tensor<B, 4>,
160        offset: Tensor<B, 4>,
161        mask: Option<Tensor<B, 4>>,
162    ) -> Tensor<B, 4> {
163        let [_batch_size, _channels_in, height_in, width_in] = input.dims();
164        let padding =
165            self.padding
166                .calculate_padding_2d(height_in, width_in, &self.kernel_size, &self.stride);
167        deform_conv2d(
168            input,
169            offset,
170            self.weight.val(),
171            mask,
172            self.bias.as_ref().map(|bias| bias.val()),
173            DeformConvOptions::new(
174                self.stride,
175                padding,
176                self.dilation,
177                self.weight_groups,
178                self.offset_groups,
179            ),
180        )
181    }
182}
183
184#[cfg(test)]
185mod tests {
186    use burn::tensor::{ElementConversion, Tolerance, ops::FloatElem};
187    type FT = FloatElem<TestBackend>;
188
189    use super::*;
190    use crate::TestBackend;
191    use burn::tensor::TensorData;
192
193    #[test]
194    fn initializer_default() {
195        let device = Default::default();
196        TestBackend::seed(&device, 0);
197
198        let config = DeformConv2dConfig::new([5, 1], [5, 5]);
199        let k = (config.channels[0] * config.kernel_size[0] * config.kernel_size[1]) as f64;
200        let k = (config.offset_groups as f64 / k).sqrt().elem::<FT>();
201        let conv = config.init::<TestBackend>(&device);
202
203        conv.weight.to_data().assert_within_range(-k..k);
204    }
205
206    #[test]
207    fn initializer_zeros() {
208        let device = Default::default();
209        TestBackend::seed(&device, 0);
210
211        let config = DeformConv2dConfig::new([5, 2], [5, 5]).with_initializer(Initializer::Zeros);
212        let conv = config.init::<TestBackend>(&device);
213
214        assert_eq!(config.initializer, Initializer::Zeros);
215        conv.weight.to_data().assert_approx_eq::<FT>(
216            &TensorData::zeros::<f32, _>(conv.weight.shape()),
217            Tolerance::default(),
218        );
219    }
220
221    #[test]
222    fn initializer_fan_out() {
223        let device = Default::default();
224        TestBackend::seed(&device, 0);
225
226        let init = Initializer::KaimingUniform {
227            gain: 1.0 / 3.0f64.sqrt(),
228            fan_out_only: true, // test that fan_out is passed to `init_with()`
229        };
230
231        let config = DeformConv2dConfig::new([5, 1], [5, 5]).with_initializer(init.clone());
232        let _ = config.init::<TestBackend>(&device);
233
234        assert_eq!(config.initializer, init);
235    }
236
237    #[test]
238    fn initializer_fan_with_groups_is_valid() {
239        let device = Default::default();
240        TestBackend::seed(&device, 0);
241
242        let init = Initializer::KaimingUniform {
243            gain: 1.0 / 3.0f64.sqrt(),
244            fan_out_only: true,
245        };
246
247        let config = DeformConv2dConfig::new([4, 4], [1, 1])
248            .with_initializer(init.clone())
249            .with_weight_groups(4);
250        let _ = config.init::<TestBackend>(&device);
251
252        assert_eq!(config.initializer, init);
253    }
254
255    #[test]
256    #[should_panic = "Both channels must be divisible by the number of groups."]
257    fn channels_with_groups_is_invalid() {
258        let device = Default::default();
259        let config = DeformConv2dConfig::new([1, 4], [1, 1]).with_weight_groups(4);
260        let _ = config.init::<TestBackend>(&device);
261    }
262
263    #[test]
264    #[should_panic = "Same padding with an even kernel size is not supported"]
265    fn same_with_even_kernel_is_invalid() {
266        let device = Default::default();
267        let config = DeformConv2dConfig::new([4, 4], [2, 2]).with_padding(PaddingConfig2d::Same);
268        let _ = config.init::<TestBackend>(&device);
269    }
270
271    #[test]
272    fn display() {
273        let config = DeformConv2dConfig::new([5, 1], [5, 5]);
274        let conv = config.init::<TestBackend>(&Default::default());
275
276        assert_eq!(
277            alloc::format!("{conv}"),
278            "DeformConv2d {stride: [1, 1], kernel_size: [5, 5], dilation: [1, 1], weight_groups: 1, offset_groups: 1, padding: Valid, params: 126}"
279        );
280    }
281
282    #[test]
283    #[should_panic = "Number of channels in input tensor and input channels of convolution must be equal. got: 4, expected: 5"]
284    fn input_channels_mismatch() {
285        let config = DeformConv2dConfig::new([5, 3], [3, 3]);
286        let conv = config.init::<TestBackend>(&Default::default());
287
288        let input = Tensor::<TestBackend, 4>::zeros([1, 4, 10, 10], &Default::default());
289        let offset = Tensor::<TestBackend, 4>::zeros([1, 2 * 3 * 3, 10, 10], &Default::default());
290        let _ = conv.forward(input, offset, None);
291    }
292}
burn_nn/modules/conv/deform_conv2d.rs

burn_nn/modules/conv/
deform_conv2d.rs
