1use crate::internal::*;
2
3use crate::ops::cnn::{PaddingSpec, Patch, PatchSpec};
4use crate::ops::nn::{BaseDataShape, DataFormat, DataShape, SymDataShape};
5
6use super::padding::ComputedPaddedDim;
7
8#[derive(Debug, Clone, new, Default, Hash, PartialEq, Eq)]
9pub struct PoolSpec {
10 pub data_format: DataFormat,
11 pub kernel_shape: TVec<usize>,
12 pub padding: PaddingSpec,
13 pub dilations: Option<TVec<usize>>,
14 pub strides: Option<TVec<usize>>,
15 pub input_channels: usize,
16 pub output_channels: usize,
17}
18
19impl PoolSpec {
20 pub fn info(&self) -> Vec<String> {
21 vec![
22 format!("Data format: {:?}", self.data_format),
23 format!(
24 "Kernel shape:{:?} (strides:{:?}, padding:{:?}, dilations:{:?})",
25 self.kernel_shape, self.strides, self.padding, self.dilations,
26 ),
27 ]
28 }
29
30 pub fn rank(&self) -> usize {
31 self.kernel_shape.len()
32 }
33
34 pub fn dilation(&self, geo_axis: usize) -> usize {
35 self.dilations.as_ref().map(|d| d[geo_axis]).unwrap_or(1)
36 }
37
38 pub fn dilations(&self) -> Cow<'_, [usize]> {
39 self.dilations
40 .as_deref()
41 .map_or_else(|| vec![1; self.kernel_shape.len()].into(), |d| d.into())
42 }
43
44 pub fn stride(&self, geo_axis: usize) -> usize {
45 self.strides.as_ref().map(|s| s[geo_axis]).unwrap_or(1)
46 }
47
48 pub fn strides(&self) -> Cow<'_, [usize]> {
49 self.strides
50 .as_deref()
51 .map_or_else(|| vec![1; self.kernel_shape.len()].into(), |d| d.into())
52 }
53
54 pub fn computed_padding<D: DimLike>(&self, input_hw: &[D]) -> TVec<ComputedPaddedDim<D>> {
55 self.padding.compute(input_hw, &self.kernel_shape, &self.dilations(), &self.strides())
56 }
57
58 pub fn output_shape<D: DimLike>(&self, input: &[D]) -> TractResult<BaseDataShape<D, TVec<D>>> {
59 let ishape: BaseDataShape<D, TVec<D>> = self.data_format.shape(input.into())?;
60 ensure!(ishape.c().to_dim() == self.input_channels.to_dim());
61 let computed = self.computed_padding(ishape.hw_dims());
62 let spatial_dims = computed.into_iter().map(|d| d.convoluted).collect::<TVec<D>>();
63 let oshape = self.data_format.from_n_c_hw(
64 ishape.n().cloned().unwrap_or_else(|| 1.into()),
65 self.output_channels.into(),
66 spatial_dims,
67 )?;
68 Ok(oshape)
69 }
70
71 pub fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
72 let oshape = self.output_shape(&inputs[0].shape)?;
73 Ok(tvec!(inputs[0].datum_type.fact(oshape.shape)))
74 }
75
76 pub fn dispose_n_axis(&self) -> PoolSpec {
77 PoolSpec { data_format: self.data_format.dispose_n_axis(), ..self.clone() }
78 }
79
80 pub fn with_n_axis(&self) -> PoolSpec {
81 PoolSpec { data_format: self.data_format.with_n(), ..self.clone() }
82 }
83
84 pub fn compute_geo(&self, input_full_shape: &[TDim]) -> TractResult<PoolGeometry> {
85 let output_shape = self.output_shape(input_full_shape)?;
86 let input_shape: SymDataShape = self.data_format.shape(input_full_shape.into())?;
87 Ok(PoolGeometry::Symbolic(SymbolicPoolGeometry {
88 pool_spec: self.clone(),
89 input_shape,
90 output_shape,
91 }))
92 }
93
94 pub fn change_geo_axes(&self, op: &AxisOp) -> TractResult<PoolSpec> {
95 let mut dilations = self.dilations().into_owned().into();
96 op.change_shape_array(&mut dilations, false)?;
97 let mut kernel_shape = self.kernel_shape.clone();
98 op.change_shape_array(&mut kernel_shape, false)?;
99 let mut strides = self.strides().into_owned().into();
100 op.change_shape_array(&mut strides, false)?;
101 let padding = self.padding.change_geo_axes(op)?;
102 Ok(PoolSpec {
103 kernel_shape,
104 padding,
105 dilations: Some(dilations),
106 strides: Some(strides),
107 ..self.clone()
108 })
109 }
110
111 pub fn declutter(&self, input: &[TDim]) -> TractResult<Option<PoolSpec>> {
112 if let PaddingSpec::ExplicitOnnxPool(before, after, _) = &self.padding {
113 let input = self.data_format.shape(input)?;
114 let input_hw = input.hw_dims();
115 let reference = self.computed_padding(input_hw);
116 for replacement in [
117 PaddingSpec::Valid,
118 PaddingSpec::SameUpper,
119 PaddingSpec::SameLower,
120 PaddingSpec::Explicit(before.clone(), after.clone()),
121 ] {
122 let new_pool_spec = PoolSpec { padding: replacement, ..self.clone() };
123 if new_pool_spec.computed_padding(input_hw) == reference {
124 return Ok(Some(new_pool_spec));
125 }
126 }
127 }
128 Ok(None)
129 }
130}
131
132pub type PoolGeometry = super::GeometryBound<SymbolicPoolGeometry, ConcretePoolGeometry>;
133
134#[derive(Debug, Clone, Hash, PartialEq, Eq)]
135pub struct SymbolicPoolGeometry {
136 pub pool_spec: PoolSpec,
137 pub input_shape: SymDataShape,
138 pub output_shape: SymDataShape,
139}
140
141#[derive(Debug, Clone, Hash, PartialEq, Eq)]
142pub struct ConcretePoolGeometry {
143 pub input_shape: DataShape,
144 pub patch: Patch,
145 pub output_shape: DataShape,
146}
147
148impl super::ResolveTo<ConcretePoolGeometry> for SymbolicPoolGeometry {
149 type Param = [usize];
150 fn resolve(&self, input_full_shape: &[usize]) -> TractResult<ConcretePoolGeometry> {
151 let input_shape = self.pool_spec.data_format.shape(input_full_shape.into())?;
152 let output_inner_stride = match self.pool_spec.data_format {
153 DataFormat::NCHW | DataFormat::CHW => 1,
154 DataFormat::NHWC | DataFormat::HWC => self.pool_spec.output_channels,
155 };
156 let mut spec = PatchSpec::for_full_shape(self.pool_spec.data_format, input_full_shape)?
157 .with_output_inner_stride(output_inner_stride)
158 .with_kernel_shape(self.pool_spec.kernel_shape.clone())
159 .with_padding(self.pool_spec.padding.clone());
160 if let Some(strides) = self.pool_spec.strides.clone() {
161 spec = spec.with_strides(strides);
162 }
163 if let Some(dilations) = self.pool_spec.dilations.clone() {
164 spec = spec.with_dilations(dilations);
165 }
166 let patch = spec.into_patch();
167 let output_shape = input_shape.fmt.from_n_c_hw(
168 *input_shape.n().unwrap_or(&1),
169 self.pool_spec.output_channels,
170 &*patch.output_shape,
171 )?;
172 Ok(ConcretePoolGeometry { input_shape, patch, output_shape })
173 }
174}