burn_nn/modules/conv/
conv3d.rs

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

burn_nn/modules/conv/
conv3d.rs
