burn_core/nn/conv/
conv_transpose1d.rs

1use alloc::format;
2
3use crate as burn;
4
5use crate::config::Config;
6use crate::module::Content;
7use crate::module::DisplaySettings;
8use crate::module::Module;
9use crate::module::ModuleDisplay;
10use crate::module::Param;
11use crate::nn::conv::checks;
12use crate::nn::Initializer;
13use crate::tensor::backend::Backend;
14use crate::tensor::module::conv_transpose1d;
15use crate::tensor::ops::ConvTransposeOptions;
16use crate::tensor::Tensor;
17
18/// Configuration to create an [1D transposed convolution](ConvTranspose1d) layer
19/// using the [init function](ConvTranspose1dConfig::init).
20#[derive(Config, Debug)]
21pub struct ConvTranspose1dConfig {
22    /// The number of channels.
23    pub channels: [usize; 2],
24    /// The size of the kernel.
25    pub kernel_size: usize,
26    /// The stride of the convolution.
27    #[config(default = "1")]
28    pub stride: usize,
29    /// Spacing between kernel elements.
30    #[config(default = "1")]
31    pub dilation: usize,
32    /// Controls the connections between input and output channels.
33    #[config(default = "1")]
34    pub groups: usize,
35    /// The padding configuration.
36    #[config(default = "0")]
37    pub padding: usize,
38    /// The padding output configuration.
39    #[config(default = "0")]
40    pub padding_out: usize,
41    /// If bias should be added to the output.
42    #[config(default = true)]
43    pub bias: bool,
44    /// The type of function used to initialize neural network parameters
45    #[config(
46        default = "Initializer::KaimingUniform{gain:1.0/num_traits::Float::sqrt(3.0),fan_out_only:false}"
47    )]
48    pub initializer: Initializer,
49}
50
51/// Applies a 1D transposed convolution over input tensors.
52#[derive(Module, Debug)]
53#[module(custom_display)]
54pub struct ConvTranspose1d<B: Backend> {
55    /// Tensor of shape `[channels_in, channels_out / groups, kernel_size]`
56    pub weight: Param<Tensor<B, 3>>,
57    /// Tensor of shape `[channels_out]`
58    pub bias: Option<Param<Tensor<B, 1>>>,
59    /// Stride of the convolution.
60    pub stride: usize,
61    /// Size of the kernel.
62    pub kernel_size: usize,
63    /// Spacing between kernel elements.
64    pub dilation: usize,
65    /// Controls the connections between input and output channels.
66    pub groups: usize,
67    /// The padding configuration.
68    pub padding: usize,
69    /// The padding output configuration.
70    pub padding_out: usize,
71    /// The number of channels.
72    pub channels: [usize; 2],
73}
74
75impl<B: Backend> ModuleDisplay for ConvTranspose1d<B> {
76    fn custom_settings(&self) -> Option<DisplaySettings> {
77        DisplaySettings::new()
78            .with_new_line_after_attribute(false)
79            .optional()
80    }
81
82    fn custom_content(&self, content: Content) -> Option<Content> {
83        content
84            .add("channels", &format!("{:?}", &self.channels))
85            .add("stride", &self.stride)
86            .add("kernel_size", &self.kernel_size)
87            .add("dilation", &self.dilation)
88            .add("groups", &self.groups)
89            .add("padding", &self.padding)
90            .add("padding_out", &self.padding_out)
91            .optional()
92    }
93}
94
95impl ConvTranspose1dConfig {
96    /// Initialize a new [conv transpose 1d](ConvTranspose1d) module.
97    pub fn init<B: Backend>(&self, device: &B::Device) -> ConvTranspose1d<B> {
98        checks::checks_channels_div_groups(self.channels[0], self.channels[1], self.groups);
99
100        let shape = [
101            self.channels[0],
102            self.channels[1] / self.groups,
103            self.kernel_size,
104        ];
105
106        let fan_in = self.channels[1] / self.groups * self.kernel_size;
107        let weight = self
108            .initializer
109            .init_with(shape, Some(fan_in), None, device);
110        let mut bias = None;
111
112        if self.bias {
113            bias = Some(
114                self.initializer
115                    .init_with([self.channels[1]], Some(fan_in), None, device),
116            );
117        }
118
119        ConvTranspose1d {
120            weight,
121            bias,
122            stride: self.stride,
123            kernel_size: self.kernel_size,
124            dilation: self.dilation,
125            groups: self.groups,
126            padding: self.padding,
127            padding_out: self.padding_out,
128            channels: self.channels,
129        }
130    }
131}
132
133impl<B: Backend> ConvTranspose1d<B> {
134    /// Applies the forward pass on the input tensor.
135    ///
136    /// See also [conv_transpose1d](crate::tensor::module::conv_transpose1d).
137    ///
138    /// # Shapes
139    ///
140    /// - input: `[batch_size, channels_in, length_in]`
141    /// - output: `[batch_size, channels_out, length_out]`
142    pub fn forward(&self, input: Tensor<B, 3>) -> Tensor<B, 3> {
143        conv_transpose1d(
144            input,
145            self.weight.val(),
146            self.bias.as_ref().map(|bias| bias.val()),
147            ConvTransposeOptions::new(
148                [self.stride],
149                [self.padding],
150                [self.padding_out],
151                [self.dilation],
152                self.groups,
153            ),
154        )
155    }
156}
157
158#[cfg(test)]
159mod tests {
160    use super::*;
161    use crate::tensor::TensorData;
162    use crate::TestBackend;
163
164    #[test]
165    fn initializer_default() {
166        TestBackend::seed(0);
167
168        let config = ConvTranspose1dConfig::new([5, 1], 5);
169        let k = (config.channels[1] * config.kernel_size) as f64;
170        let k = (config.groups as f64 / k).sqrt() as f32;
171        let conv = config.init::<TestBackend>(&Default::default());
172
173        conv.weight.to_data().assert_within_range(-k..k);
174    }
175
176    #[test]
177    fn initializer_zeros() {
178        TestBackend::seed(0);
179
180        let config = ConvTranspose1dConfig::new([5, 2], 5).with_initializer(Initializer::Zeros);
181        let conv = config.init::<TestBackend>(&Default::default());
182
183        assert_eq!(config.initializer, Initializer::Zeros);
184        conv.weight
185            .to_data()
186            .assert_approx_eq(&TensorData::zeros::<f32, _>(conv.weight.shape()), 3);
187    }
188
189    #[test]
190    fn display() {
191        let config = ConvTranspose1dConfig::new([5, 2], 5);
192        let conv = config.init::<TestBackend>(&Default::default());
193
194        assert_eq!(
195            format!("{}", conv),
196            "ConvTranspose1d {channels: [5, 2], stride: 1, kernel_size: 5, dilation: 1, groups: 1, padding: 0, padding_out: 0, params: 52}"
197        );
198    }
199}
burn_core/nn/conv/conv_transpose1d.rs

burn_core/nn/conv/
conv_transpose1d.rs
