burn_core/nn/conv/
conv2d.rs

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

burn_core/nn/conv/
conv2d.rs
