burn_nn/modules/conv/
conv1d.rs

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

burn_nn/modules/conv/
conv1d.rs
