burn_core/nn/conv/
deform_conv2d.rs

1use alloc::format;
2use burn_tensor::ops::DeformConvOptions;
3
4use crate as burn;
5
6use crate::config::Config;
7use crate::module::{Content, DisplaySettings, Ignored, Module, ModuleDisplay, Param};
8use crate::nn::Initializer;
9use crate::nn::PaddingConfig2d;
10use crate::tensor::backend::Backend;
11use crate::tensor::module::deform_conv2d;
12use crate::tensor::Tensor;
13
14use crate::nn::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: Ignored<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: Ignored(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        // Since padding does not implement ModuleDisplay, we need to format it manually.
131        let padding_formatted = format!("{}", &self.padding);
132
133        // Format the stride, kernel_size and dilation as strings, formatted as arrays instead of indexed.
134        let stride = format!("{:?}", self.stride);
135        let kernel_size = format!("{:?}", self.kernel_size);
136        let dilation = format!("{:?}", self.dilation);
137
138        content
139            .add("stride", &stride)
140            .add("kernel_size", &kernel_size)
141            .add("dilation", &dilation)
142            .add("weight_groups", &self.weight_groups)
143            .add("offset_groups", &self.offset_groups)
144            .add("padding", &padding_formatted)
145            .optional()
146    }
147}
148
149impl<B: Backend> DeformConv2d<B> {
150    /// Applies the forward pass on the input tensor.
151    ///
152    /// See [deform_conv2d](crate::tensor::module::deform_conv2d) for more information.
153    ///
154    /// # Shapes
155    ///
156    /// - input: `[batch_size, channels_in, height_in, width_in]`
157    /// - offset: `[batch_size, 2 * offset_groups * kernel_height * kernel_width, height_out, width_out]`
158    /// - mask: `[batch_size, offset_groups * kernel_height * kernel_width, height_out, width_out]`
159    /// - output: `[batch_size, channels_out, height_out, width_out]`
160    pub fn forward(
161        &self,
162        input: Tensor<B, 4>,
163        offset: Tensor<B, 4>,
164        mask: Option<Tensor<B, 4>>,
165    ) -> Tensor<B, 4> {
166        let [_batch_size, _channels_in, height_in, width_in] = input.dims();
167        let padding =
168            self.padding
169                .calculate_padding_2d(height_in, width_in, &self.kernel_size, &self.stride);
170        deform_conv2d(
171            input,
172            offset,
173            self.weight.val(),
174            mask,
175            self.bias.as_ref().map(|bias| bias.val()),
176            DeformConvOptions::new(
177                self.stride,
178                padding,
179                self.dilation,
180                self.weight_groups,
181                self.offset_groups,
182            ),
183        )
184    }
185}
186
187#[cfg(test)]
188mod tests {
189    use super::*;
190    use crate::tensor::TensorData;
191    use crate::TestBackend;
192
193    #[test]
194    fn initializer_default() {
195        TestBackend::seed(0);
196
197        let config = DeformConv2dConfig::new([5, 1], [5, 5]);
198        let k = (config.channels[0] * config.kernel_size[0] * config.kernel_size[1]) as f64;
199        let k = (config.offset_groups as f64 / k).sqrt() as f32;
200        let device = Default::default();
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        TestBackend::seed(0);
209
210        let config = DeformConv2dConfig::new([5, 2], [5, 5]).with_initializer(Initializer::Zeros);
211        let device = Default::default();
212        let conv = config.init::<TestBackend>(&device);
213
214        assert_eq!(config.initializer, Initializer::Zeros);
215        conv.weight
216            .to_data()
217            .assert_approx_eq(&TensorData::zeros::<f32, _>(conv.weight.shape()), 3);
218    }
219
220    #[test]
221    fn initializer_fan_out() {
222        TestBackend::seed(0);
223
224        let init = Initializer::KaimingUniform {
225            gain: 1.0 / 3.0f64.sqrt(),
226            fan_out_only: true, // test that fan_out is passed to `init_with()`
227        };
228        let device = Default::default();
229        let config = DeformConv2dConfig::new([5, 1], [5, 5]).with_initializer(init.clone());
230        let _ = config.init::<TestBackend>(&device);
231
232        assert_eq!(config.initializer, init);
233    }
234
235    #[test]
236    fn initializer_fan_with_groups_is_valid() {
237        TestBackend::seed(0);
238
239        let init = Initializer::KaimingUniform {
240            gain: 1.0 / 3.0f64.sqrt(),
241            fan_out_only: true,
242        };
243        let device = Default::default();
244        let config = DeformConv2dConfig::new([4, 4], [1, 1])
245            .with_initializer(init.clone())
246            .with_weight_groups(4);
247        let _ = config.init::<TestBackend>(&device);
248
249        assert_eq!(config.initializer, init);
250    }
251
252    #[test]
253    #[should_panic = "Both channels must be divisible by the number of groups."]
254    fn channels_with_groups_is_invalid() {
255        let device = Default::default();
256        let config = DeformConv2dConfig::new([1, 4], [1, 1]).with_weight_groups(4);
257        let _ = config.init::<TestBackend>(&device);
258    }
259
260    #[test]
261    #[should_panic = "Same padding with an even kernel size is not supported"]
262    fn same_with_even_kernel_is_invalid() {
263        let device = Default::default();
264        let config = DeformConv2dConfig::new([4, 4], [2, 2]).with_padding(PaddingConfig2d::Same);
265        let _ = config.init::<TestBackend>(&device);
266    }
267
268    #[test]
269    fn display() {
270        let config = DeformConv2dConfig::new([5, 1], [5, 5]);
271        let conv = config.init::<TestBackend>(&Default::default());
272
273        assert_eq!(
274            alloc::format!("{}", conv),
275            "DeformConv2d {stride: [1, 1], kernel_size: [5, 5], dilation: [1, 1], weight_groups: 1, offset_groups: 1, padding: Valid, params: 126}"
276        );
277    }
278}
burn_core/nn/conv/deform_conv2d.rs

burn_core/nn/conv/
deform_conv2d.rs
