diff --git a/crates/burn-core/src/nn/conv/checks.rs b/crates/burn-core/src/nn/conv/checks.rs index cd346163ad..36932621f1 100644 --- a/crates/burn-core/src/nn/conv/checks.rs +++ b/crates/burn-core/src/nn/conv/checks.rs @@ -9,3 +9,14 @@ pub(crate) fn checks_channels_div_groups(channels_in: usize, channels_out: usize ); } } + +// https://github.com/tracel-ai/burn/issues/2676 +/// Only symmetric padding is currently supported. As such, using `Same` padding with an even kernel +/// size is not supported as it will not produce the same output size. +pub(crate) fn check_same_padding_support(kernel_size: &[usize]) { + for k in kernel_size.iter() { + if k % 2 == 0 { + unimplemented!("Same padding with an even kernel size is not supported"); + } + } +} diff --git a/crates/burn-core/src/nn/conv/conv1d.rs b/crates/burn-core/src/nn/conv/conv1d.rs index 0b64eab324..c3f61a6b07 100644 --- a/crates/burn-core/src/nn/conv/conv1d.rs +++ b/crates/burn-core/src/nn/conv/conv1d.rs @@ -28,6 +28,10 @@ pub struct Conv1dConfig { #[config(default = "1")] pub groups: usize, /// The padding configuration. + /// + /// ### Warning + /// Only symmetric padding is currently supported. As such, using `Same` padding with an even kernel + /// size is not supported as it will not produce the same output size. #[config(default = "PaddingConfig1d::Valid")] pub padding: PaddingConfig1d, /// If bias should be added to the output. @@ -87,6 +91,9 @@ impl Conv1dConfig { /// Initialize a new [conv1d](Conv1d) module. pub fn init(&self, device: &B::Device) -> Conv1d { checks::checks_channels_div_groups(self.channels_in, self.channels_out, self.groups); + if self.padding == PaddingConfig1d::Same { + checks::check_same_padding_support(&[self.kernel_size]); + } let shape = [ self.channels_out, @@ -175,6 +182,14 @@ mod tests { .assert_approx_eq(&TensorData::zeros::(conv.weight.shape()), 3); } + #[test] + #[should_panic = "Same padding with an even kernel size is not supported"] + fn same_with_even_kernel_is_invalid() { + let device = Default::default(); + let config = Conv1dConfig::new(5, 5, 4).with_padding(PaddingConfig1d::Same); + let _ = config.init::(&device); + } + #[test] fn display() { let config = Conv1dConfig::new(5, 5, 5); diff --git a/crates/burn-core/src/nn/conv/conv2d.rs b/crates/burn-core/src/nn/conv/conv2d.rs index 73be36d357..72c00187be 100644 --- a/crates/burn-core/src/nn/conv/conv2d.rs +++ b/crates/burn-core/src/nn/conv/conv2d.rs @@ -30,6 +30,10 @@ pub struct Conv2dConfig { #[config(default = "1")] pub groups: usize, /// The padding configuration. + /// + /// ### Warning + /// Only symmetric padding is currently supported. As such, using `Same` padding with an even kernel + /// size is not supported as it will not produce the same output size. #[config(default = "PaddingConfig2d::Valid")] pub padding: PaddingConfig2d, /// If bias should be added to the output. @@ -68,6 +72,9 @@ impl Conv2dConfig { /// Initialize a new [conv2d](Conv2d) module. pub fn init(&self, device: &B::Device) -> Conv2d { checks::checks_channels_div_groups(self.channels[0], self.channels[1], self.groups); + if self.padding == PaddingConfig2d::Same { + checks::check_same_padding_support(&self.kernel_size); + } let shape = [ self.channels[1], @@ -228,6 +235,14 @@ mod tests { let _ = config.init::(&device); } + #[test] + #[should_panic = "Same padding with an even kernel size is not supported"] + fn same_with_even_kernel_is_invalid() { + let device = Default::default(); + let config = Conv2dConfig::new([4, 4], [2, 2]).with_padding(PaddingConfig2d::Same); + let _ = config.init::(&device); + } + #[test] fn display() { let config = Conv2dConfig::new([5, 1], [5, 5]); diff --git a/crates/burn-core/src/nn/conv/conv3d.rs b/crates/burn-core/src/nn/conv/conv3d.rs index de7fb1ce2b..0b5d530c5a 100644 --- a/crates/burn-core/src/nn/conv/conv3d.rs +++ b/crates/burn-core/src/nn/conv/conv3d.rs @@ -68,6 +68,9 @@ impl Conv3dConfig { /// Initialize a new [conv3d](Conv3d) module. pub fn init(&self, device: &B::Device) -> Conv3d { checks::checks_channels_div_groups(self.channels[0], self.channels[1], self.groups); + if self.padding == PaddingConfig3d::Same { + checks::check_same_padding_support(&self.kernel_size); + } let shape = [ self.channels[1], @@ -228,6 +231,14 @@ mod tests { assert_eq!(config.initializer, init); } + #[test] + #[should_panic = "Same padding with an even kernel size is not supported"] + fn same_with_even_kernel_is_invalid() { + let device = Default::default(); + let config = Conv3dConfig::new([4, 4], [2, 2, 2]).with_padding(PaddingConfig3d::Same); + let _ = config.init::(&device); + } + #[test] fn display() { let config = Conv3dConfig::new([5, 1], [5, 5, 5]); diff --git a/crates/burn-core/src/nn/conv/deform_conv2d.rs b/crates/burn-core/src/nn/conv/deform_conv2d.rs index 03becd9d4e..2baff11d07 100644 --- a/crates/burn-core/src/nn/conv/deform_conv2d.rs +++ b/crates/burn-core/src/nn/conv/deform_conv2d.rs @@ -33,6 +33,10 @@ pub struct DeformConv2dConfig { #[config(default = "1")] pub offset_groups: usize, /// The padding configuration. + /// + /// ### Warning + /// Only symmetric padding is currently supported. As such, using `Same` padding with an even kernel + /// size is not supported as it will not produce the same output size. #[config(default = "PaddingConfig2d::Valid")] pub padding: PaddingConfig2d, /// If bias should be added to the output. @@ -73,6 +77,9 @@ impl DeformConv2dConfig { /// Initialize a new [DeformConv2d](DeformConv2d) module. pub fn init(&self, device: &B::Device) -> DeformConv2d { checks::checks_channels_div_groups(self.channels[0], self.channels[1], self.weight_groups); + if self.padding == PaddingConfig2d::Same { + checks::check_same_padding_support(&self.kernel_size); + } let shape = [ self.channels[1], @@ -250,6 +257,14 @@ mod tests { let _ = config.init::(&device); } + #[test] + #[should_panic = "Same padding with an even kernel size is not supported"] + fn same_with_even_kernel_is_invalid() { + let device = Default::default(); + let config = DeformConv2dConfig::new([4, 4], [2, 2]).with_padding(PaddingConfig2d::Same); + let _ = config.init::(&device); + } + #[test] fn display() { let config = DeformConv2dConfig::new([5, 1], [5, 5]); diff --git a/crates/burn-core/src/nn/pool/avg_pool1d.rs b/crates/burn-core/src/nn/pool/avg_pool1d.rs index 949160fd5b..24ec8ff972 100644 --- a/crates/burn-core/src/nn/pool/avg_pool1d.rs +++ b/crates/burn-core/src/nn/pool/avg_pool1d.rs @@ -1,4 +1,5 @@ use crate as burn; +use crate::nn::conv::checks::check_same_padding_support; use crate::config::Config; use crate::module::{Content, DisplaySettings, ModuleDisplay}; @@ -18,6 +19,10 @@ pub struct AvgPool1dConfig { #[config(default = "1")] pub stride: usize, /// The padding configuration. + /// + /// ### Warning + /// Only symmetric padding is currently supported. As such, using `Same` padding with an even kernel + /// size is not supported as it will not produce the same output size. #[config(default = "PaddingConfig1d::Valid")] pub padding: PaddingConfig1d, /// If the padding is counted in the denominator when computing the average. @@ -36,10 +41,6 @@ pub struct AvgPool1dConfig { /// legitimate values, and they contribute to the denominator /// when calculating the average. This is equivalent to /// `torch.nn.AvgPool2d` with `count_include_pad=True`. -/// -/// TODO: Add support for `count_include_pad=False`, see -/// [Issue 636](https://github.com/tracel-ai/burn/issues/636) - #[derive(Module, Clone, Debug)] #[module(custom_display)] pub struct AvgPool1d { @@ -73,6 +74,9 @@ impl ModuleDisplay for AvgPool1d { impl AvgPool1dConfig { /// Initialize a new [avg pool 1d](AvgPool1d) module. pub fn init(&self) -> AvgPool1d { + if self.padding == PaddingConfig1d::Same { + check_same_padding_support(&[self.kernel_size]); + } AvgPool1d { stride: self.stride, kernel_size: self.kernel_size, @@ -111,6 +115,13 @@ impl AvgPool1d { mod tests { use super::*; + #[test] + #[should_panic = "Same padding with an even kernel size is not supported"] + fn same_with_even_kernel_is_invalid() { + let config = AvgPool1dConfig::new(2).with_padding(PaddingConfig1d::Same); + let _ = config.init(); + } + #[test] fn display() { let config = AvgPool1dConfig::new(3); diff --git a/crates/burn-core/src/nn/pool/avg_pool2d.rs b/crates/burn-core/src/nn/pool/avg_pool2d.rs index 6c6ffc87ed..343d59922b 100644 --- a/crates/burn-core/src/nn/pool/avg_pool2d.rs +++ b/crates/burn-core/src/nn/pool/avg_pool2d.rs @@ -1,4 +1,5 @@ use crate as burn; +use crate::nn::conv::checks::check_same_padding_support; use crate::config::Config; use crate::module::{Content, DisplaySettings, ModuleDisplay}; @@ -18,6 +19,10 @@ pub struct AvgPool2dConfig { #[config(default = "[1, 1]")] pub strides: [usize; 2], /// The padding configuration. + /// + /// ### Warning + /// Only symmetric padding is currently supported. As such, using `Same` padding with an even kernel + /// size is not supported as it will not produce the same output size. #[config(default = "PaddingConfig2d::Valid")] pub padding: PaddingConfig2d, /// If the padding is counted in the denominator when computing the average. @@ -36,9 +41,6 @@ pub struct AvgPool2dConfig { /// legitimate values, and they contribute to the denominator /// when calculating the average. This is equivalent to /// `torch.nn.AvgPool2d` with `count_include_pad=True`. -/// -/// TODO: Add support for `count_include_pad=False`, see -/// [Issue 636](https://github.com/tracel-ai/burn/issues/636) #[derive(Module, Clone, Debug)] #[module(custom_display)] pub struct AvgPool2d { @@ -72,6 +74,9 @@ impl ModuleDisplay for AvgPool2d { impl AvgPool2dConfig { /// Initialize a new [avg pool 2d](AvgPool2d) module. pub fn init(&self) -> AvgPool2d { + if self.padding == PaddingConfig2d::Same { + check_same_padding_support(&self.kernel_size); + } AvgPool2d { stride: self.strides, kernel_size: self.kernel_size, @@ -110,6 +115,13 @@ impl AvgPool2d { mod tests { use super::*; + #[test] + #[should_panic = "Same padding with an even kernel size is not supported"] + fn same_with_even_kernel_is_invalid() { + let config = AvgPool2dConfig::new([2, 2]).with_padding(PaddingConfig2d::Same); + let _ = config.init(); + } + #[test] fn display() { let config = AvgPool2dConfig::new([3, 3]); diff --git a/crates/burn-core/src/nn/pool/max_pool1d.rs b/crates/burn-core/src/nn/pool/max_pool1d.rs index 5be363e908..71041e6155 100644 --- a/crates/burn-core/src/nn/pool/max_pool1d.rs +++ b/crates/burn-core/src/nn/pool/max_pool1d.rs @@ -1,4 +1,5 @@ use crate as burn; +use crate::nn::conv::checks::check_same_padding_support; use crate::config::Config; use crate::module::{Content, DisplaySettings, ModuleDisplay}; @@ -18,6 +19,10 @@ pub struct MaxPool1dConfig { #[config(default = "1")] pub stride: usize, /// The padding configuration. + /// + /// ### Warning + /// Only symmetric padding is currently supported. As such, using `Same` padding with an even kernel + /// size is not supported as it will not produce the same output size. #[config(default = "PaddingConfig1d::Valid")] pub padding: PaddingConfig1d, /// The dilation. @@ -61,6 +66,9 @@ impl ModuleDisplay for MaxPool1d { impl MaxPool1dConfig { /// Initialize a new [max pool 1d](MaxPool1d) module. pub fn init(&self) -> MaxPool1d { + if self.padding == PaddingConfig1d::Same { + check_same_padding_support(&[self.kernel_size]); + } MaxPool1d { stride: self.stride, kernel_size: self.kernel_size, @@ -93,6 +101,13 @@ impl MaxPool1d { mod tests { use super::*; + #[test] + #[should_panic = "Same padding with an even kernel size is not supported"] + fn same_with_even_kernel_is_invalid() { + let config = MaxPool1dConfig::new(2).with_padding(PaddingConfig1d::Same); + let _ = config.init(); + } + #[test] fn display() { let config = MaxPool1dConfig::new(3); diff --git a/crates/burn-core/src/nn/pool/max_pool2d.rs b/crates/burn-core/src/nn/pool/max_pool2d.rs index ab9c60d276..3eb94f5db5 100644 --- a/crates/burn-core/src/nn/pool/max_pool2d.rs +++ b/crates/burn-core/src/nn/pool/max_pool2d.rs @@ -1,4 +1,5 @@ use crate as burn; +use crate::nn::conv::checks::check_same_padding_support; use crate::config::Config; use crate::module::{Content, DisplaySettings, ModuleDisplay}; @@ -18,6 +19,10 @@ pub struct MaxPool2dConfig { #[config(default = "[1, 1]")] pub strides: [usize; 2], /// The padding configuration. + /// + /// ### Warning + /// Only symmetric padding is currently supported. As such, using `Same` padding with an even kernel + /// size is not supported as it will not produce the same output size. #[config(default = "PaddingConfig2d::Valid")] pub padding: PaddingConfig2d, /// The dilation. @@ -61,6 +66,9 @@ impl ModuleDisplay for MaxPool2d { impl MaxPool2dConfig { /// Initialize a new [max pool 2d](MaxPool2d) module. pub fn init(&self) -> MaxPool2d { + if self.padding == PaddingConfig2d::Same { + check_same_padding_support(&self.kernel_size); + } MaxPool2d { stride: self.strides, kernel_size: self.kernel_size, @@ -93,6 +101,13 @@ impl MaxPool2d { mod tests { use super::*; + #[test] + #[should_panic = "Same padding with an even kernel size is not supported"] + fn same_with_even_kernel_is_invalid() { + let config = MaxPool2dConfig::new([2, 2]).with_padding(PaddingConfig2d::Same); + let _ = config.init(); + } + #[test] fn display() { let config = MaxPool2dConfig::new([3, 3]);