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::Tensor;
11use crate::tensor::backend::Backend;
12use crate::tensor::module::deform_conv2d;
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 burn_tensor::Tolerance;
190
191    use super::*;
192    use crate::TestBackend;
193    use crate::tensor::TensorData;
194
195    #[test]
196    fn initializer_default() {
197        TestBackend::seed(0);
198
199        let config = DeformConv2dConfig::new([5, 1], [5, 5]);
200        let k = (config.channels[0] * config.kernel_size[0] * config.kernel_size[1]) as f64;
201        let k = (config.offset_groups as f64 / k).sqrt() as f32;
202        let device = Default::default();
203        let conv = config.init::<TestBackend>(&device);
204
205        conv.weight.to_data().assert_within_range(-k..k);
206    }
207
208    #[test]
209    fn initializer_zeros() {
210        TestBackend::seed(0);
211
212        let config = DeformConv2dConfig::new([5, 2], [5, 5]).with_initializer(Initializer::Zeros);
213        let device = Default::default();
214        let conv = config.init::<TestBackend>(&device);
215
216        assert_eq!(config.initializer, Initializer::Zeros);
217        conv.weight.to_data().assert_approx_eq::<f32>(
218            &TensorData::zeros::<f32, _>(conv.weight.shape()),
219            Tolerance::default(),
220        );
221    }
222
223    #[test]
224    fn initializer_fan_out() {
225        TestBackend::seed(0);
226
227        let init = Initializer::KaimingUniform {
228            gain: 1.0 / 3.0f64.sqrt(),
229            fan_out_only: true, // test that fan_out is passed to `init_with()`
230        };
231        let device = Default::default();
232        let config = DeformConv2dConfig::new([5, 1], [5, 5]).with_initializer(init.clone());
233        let _ = config.init::<TestBackend>(&device);
234
235        assert_eq!(config.initializer, init);
236    }
237
238    #[test]
239    fn initializer_fan_with_groups_is_valid() {
240        TestBackend::seed(0);
241
242        let init = Initializer::KaimingUniform {
243            gain: 1.0 / 3.0f64.sqrt(),
244            fan_out_only: true,
245        };
246        let device = Default::default();
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_core/nn/conv/deform_conv2d.rs

burn_core/nn/conv/
deform_conv2d.rs
