tract_core/ops/cnn/deconv/
deconv.rs

1use crate::internal::*;
2use crate::ops::array::MultiBroadcastTo;
3use crate::ops::cnn::KernelFormat;
4use crate::ops::cnn::PoolSpec;
5use crate::ops::cnn::wire_reshape_bias_for_bin;
6use crate::ops::einsum::EinSum;
7
8#[derive(Clone, Debug, new, Hash, PartialEq, Eq)]
9pub struct Deconv {
10    pub pool_spec: PoolSpec,
11    pub kernel_format: KernelFormat,
12    pub adjustments: TVec<usize>,
13    pub group: usize,
14}
15
16impl Deconv {
17    /// Whether every output pixel receives exactly one kernel tap whatever the input size:
18    /// on each spatial axis the kernel extent equals the stride, with no dilation, no
19    /// adjustment and no padding. `DeconvSum` then takes the bias unbroadcast.
20    fn taps_never_overlap(&self) -> bool {
21        let strides = self.pool_spec.strides();
22        let dilations = self.pool_spec.dilations();
23        (0..self.pool_spec.rank()).all(|ax| {
24            self.pool_spec.kernel_shape[ax] == strides[ax]
25                && dilations[ax] == 1
26                && self.adjustments[ax] == 0
27                && self.pool_spec.padding.valid_dim(ax, false)
28        })
29    }
30
31    fn wire_with_deconv_sum(
32        &self,
33        name: &str,
34        target: &mut TypedModel,
35        inputs: &[OutletId],
36    ) -> TractResult<TVec<OutletId>> {
37        let input_shape = target.outlet_fact(inputs[0])?.shape.clone();
38        let shape = self.pool_spec.data_format.shape(input_shape.to_tvec())?;
39        let geo_dim = shape.hw_dims().iter().product();
40
41        // collapse H and W together in input: (N) I HW or (N) HW I
42        let mut input = target.wire_node(
43            format!("{name}.reshaped_input"),
44            AxisOp::Reshape(shape.h_axis(), shape.hw_dims().into(), tvec!(geo_dim)),
45            &[inputs[0]],
46        )?;
47
48        // rework input to (N) (G) I/G HW or (N) (G) HW I/G
49        if self.group != 1 {
50            // input is (N) HW I or (N) I HW
51            let i_axis = self.pool_spec.data_format.has_n() as usize
52                + self.pool_spec.data_format.c_is_last() as usize;
53            let i_dim = target.outlet_fact(input[0])?.shape[i_axis].clone();
54            input = target.wire_node(
55                format!("{name}.reshaped_input_for_group"),
56                AxisOp::Reshape(
57                    i_axis,
58                    tvec![i_dim.clone()],
59                    tvec!(self.group.to_dim(), i_dim / self.group),
60                ),
61                &input,
62            )?;
63            if self.pool_spec.data_format.c_is_last() {
64                input = target.wire_node(
65                    format!("{name}.group_axis_left"),
66                    AxisOp::Move(
67                        self.pool_spec.data_format.has_n() as usize + 1,
68                        self.pool_spec.data_format.has_n() as usize,
69                    ),
70                    &input,
71                )?;
72            }
73        }
74
75        let mut kernel = tvec!(inputs[1]);
76        let kernel_fact = target.outlet_fact(kernel[0])?.clone();
77        for (ix, op) in self
78            .kernel_format
79            .kernel_as_group_o_i_hw_ops(&kernel_fact.shape, self.group)
80            .into_iter()
81            .enumerate()
82        {
83            kernel = target.wire_node(format!("{name}.kernel.{ix}"), op, &kernel)?;
84        }
85
86        kernel = target.wire_node(format!("{name}.kernel.mv_i"), AxisOp::Move(2, 3), &kernel)?;
87        kernel =
88            AxisOp::wire_collapse_axis(target, format!("{name}.kernel.col_ohw"), kernel[0], 1)?;
89        if self.group == 1 {
90            kernel = target.wire_node(format!("{name}.kernel.rm_g"), AxisOp::Rm(0), &kernel)?;
91        }
92        // Depthwise deconvolution: the einsum's contraction dim is I/G == 1, so it
93        // degenerates to an outer product that materialises `[N, C, HkWk, HW]` for
94        // DeconvSum to scatter -- and most of it lands out of bounds and is discarded
95        // (GTCRN: 3267 products per channel for 33 outputs). Fold the multiply into the
96        // accumulation instead.
97        // Rank 2 only: that is the path with the specialised pointer walk. The generic
98        // fused loop builds a coordinate per output element and is far slower than the
99        // einsum it would replace, so anything else keeps the existing lowering.
100        let is_depthwise = self.group > 1
101            && self.group == self.pool_spec.input_channels
102            && self.group == self.pool_spec.output_channels
103            && self.pool_spec.rank() == 2;
104        if is_depthwise
105            && !self.pool_spec.data_format.c_is_last()
106            && !super::deconv_sum::TRACT_DISABLE_DEPTHWISE_DECONV.get()
107        {
108            let mut bias = wire_reshape_bias_for_bin(
109                target,
110                format!("{name}.reshape_bias"),
111                inputs[2],
112                shape.rank(),
113                shape.c_axis(),
114                self.pool_spec.output_channels,
115            )?[0];
116            let output_shape =
117                super::output_shape(&self.pool_spec, &shape.shape, &self.adjustments)?;
118            bias = target.wire_node(
119                format!("{name}.broadcast_bias"),
120                MultiBroadcastTo { shape: output_shape.into() },
121                &[bias],
122            )?[0];
123            return target.wire_node(
124                format!("{name}.depthwise_deconv_sum"),
125                super::deconv_sum::DepthwiseDeconv::new(
126                    self.pool_spec.clone(),
127                    self.kernel_format,
128                    input_shape,
129                    self.adjustments.clone(),
130                    self.group,
131                ),
132                &[kernel[0], input[0], bias],
133            );
134        }
135
136        let mut expr = if self.pool_spec.data_format.c_is_last() {
137            "gmk,Ngnk->Ngmn".to_string()
138        } else {
139            "gmk,Ngkn->Ngmn".to_string()
140        };
141        if !self.pool_spec.data_format.has_n() {
142            expr = expr.replace('N', "");
143        }
144        if self.group == 1 {
145            expr = expr.replace('g', "");
146        }
147        let einsum = target.wire_node(
148            format!("{name}.einsum"),
149            EinSum { axes: expr.parse()?, operating_dt: kernel_fact.datum_type, q_params: None },
150            &[kernel[0], input[0]],
151        )?;
152
153        let mut bias = wire_reshape_bias_for_bin(
154            target,
155            format!("{name}.reshape_bias"),
156            inputs[2],
157            shape.rank(),
158            shape.c_axis(),
159            self.pool_spec.output_channels,
160        )?[0];
161        if !self.taps_never_overlap() {
162            let output_shape =
163                super::output_shape(&self.pool_spec, &shape.shape, &self.adjustments)?;
164            bias = target.wire_node(
165                format!("{name}.broadcast_bias"),
166                MultiBroadcastTo { shape: output_shape.into() },
167                &[bias],
168            )?[0];
169        }
170
171        // einsum must be (N_)CHkWk_HW
172        let deconv_sum = target.wire_node(
173            format!("{name}.deconv_sum"),
174            super::deconv_sum::DeconvSum::new(
175                self.pool_spec.clone(),
176                self.kernel_format,
177                input_shape,
178                self.adjustments.clone(),
179                self.group,
180            ),
181            &[einsum[0], bias],
182        )?;
183        Ok(deconv_sum)
184    }
185}
186
187impl Op for Deconv {
188    fn name(&self) -> StaticName {
189        "Deconv".into()
190    }
191
192    fn info(&self) -> TractResult<Vec<String>> {
193        Ok(vec![format!("{:?}", self.pool_spec)])
194    }
195
196    op_as_typed_op!();
197}
198
199impl EvalOp for Deconv {
200    op_out_of_plan!();
201
202    fn eval(&self, _ctx: &EvalContext, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
203        ensure!(inputs.len() == 3);
204        let mut model = TypedModel::default();
205        let inputs = inputs
206            .into_iter()
207            .enumerate()
208            .map(|(ix, input)| model.add_const(format!("s{ix}"), input.into_tensor()))
209            .collect::<TractResult<TVec<OutletId>>>()?;
210        let output = self.wire_with_deconv_sum("adhoc", &mut model, &inputs)?;
211        model.select_output_outlets(&output)?;
212        model.into_runnable()?.run(tvec![]).context("In adhoc deconvolution eval")
213    }
214}
215
216impl TypedOp for Deconv {
217    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
218        ensure!(inputs.len() == 3);
219        let x_fact = inputs[0];
220        let k_fact = inputs[1];
221        ensure!(
222            &self.pool_spec.input_channels.to_dim()
223                == self.pool_spec.data_format.shape(&inputs[0].shape)?.c()
224        );
225        ensure!(
226            self.pool_spec.input_channels.to_dim()
227                == *self.kernel_format.input_channels(&k_fact.shape, self.group)
228        );
229        let output_shape = super::output_shape(&self.pool_spec, &x_fact.shape, &self.adjustments)?;
230        Ok(tvec!(x_fact.datum_type.fact(&output_shape)))
231    }
232
233    fn axes_mapping(
234        &self,
235        inputs: &[&TypedFact],
236        outputs: &[&TypedFact],
237    ) -> TractResult<AxesMapping> {
238        let fact = &inputs[0];
239        let k_fact = &inputs[1];
240        let shape = self.pool_spec.data_format.shape(&fact.shape)?;
241        let mut axes = AxesMapping::disconnected(inputs, outputs)?
242            .renaming((InOut::In(0), shape.c_axis()), 'I')?
243            .renaming((InOut::Out(0), shape.c_axis()), 'O')?;
244        if let Some(n_axis) = shape.n_axis() {
245            axes = axes
246                .renaming((InOut::In(0), n_axis), 'N')?
247                .linking('N', (InOut::Out(0), n_axis))?;
248        }
249        let h_axis = shape.h_axis();
250        let geo = "HWXYZ".chars().chain('a'..);
251        let kernel_spatial_shape = self.kernel_format.spatial_shape(&k_fact.shape);
252        for ((ix, dim), repr) in kernel_spatial_shape.iter().enumerate().zip(geo) {
253            if dim.is_one()
254                && self.pool_spec.stride(ix) == 1
255                && self.pool_spec.padding.valid_dim(ix, true)
256                && self.adjustments[ix] == 0
257            {
258                axes = axes
259                    .renaming((InOut::In(0), ix + h_axis), repr)?
260                    .linking((InOut::In(0), ix + h_axis), (InOut::Out(0), ix + h_axis))?;
261            }
262        }
263        Ok(axes)
264    }
265
266    fn codegen(
267        &self,
268        model: &TypedModel,
269        node: &TypedNode,
270    ) -> TractResult<Option<TypedModelPatch>> {
271        let mut patch = TypedModelPatch::default();
272        let inputs = patch.taps(model, &node.inputs)?;
273        let output = self
274            .wire_with_deconv_sum(&node.name, &mut patch, &inputs)
275            .context("In wire_with_deconv_sum")?;
276        patch.shunt_outside(model, node.id.into(), output[0])?;
277        Ok(Some(patch))
278    }
279
280    as_op!();
281}
tract_core/ops/cnn/deconv/deconv.rs

tract_core/ops/cnn/deconv/
deconv.rs
