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

tract_core/ops/cnn/deconv/
deconv.rs
