tract_core/ops/downsample/
mod.rs

1use crate::internal::*;
2use crate::ops;
3use ndarray::prelude::*;
4
5use super::identity::Identity;
6
7mod array;
8mod conv;
9mod scan;
10
11#[derive(Debug, Clone, new, Default, PartialEq, Eq, Hash)]
12pub struct Downsample {
13    pub axis: usize,
14    pub stride: isize,
15    pub modulo: usize,
16}
17
18impl Downsample {
19    pub(crate) fn transform_dim(&self, input_dim: &TDim) -> TDim {
20        (input_dim.clone() - self.modulo).div_ceil(self.stride.unsigned_abs() as u64)
21    }
22
23    pub(crate) fn transform_fact(&self, input_fact: &TypedFact) -> TractResult<TypedFact> {
24        let mut downed = input_fact.clone();
25        let down_len = self.transform_dim(&input_fact.shape[self.axis]);
26        downed.shape.set(self.axis, down_len);
27        if let Some(k) = downed.konst {
28            let mut outputs = self.eval(tvec!(k.into_tvalue()))?;
29            downed.konst = Some(outputs.remove(0).into_arc_tensor())
30        }
31        if cfg!(debug_assertions) {
32            downed.consistent()?;
33        }
34        Ok(downed)
35    }
36}
37
38impl Op for Downsample {
39    fn name(&self) -> Cow<str> {
40        "Downsample".into()
41    }
42
43    fn info(&self) -> TractResult<Vec<String>> {
44        Ok(vec![format!("axis:{} stride:{} modulo:{}", self.axis, self.stride, self.modulo)])
45    }
46
47    impl_op_same_as!();
48    op_as_typed_op!();
49}
50
51impl EvalOp for Downsample {
52    fn is_stateless(&self) -> bool {
53        true
54    }
55
56    fn eval(&self, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
57        let input = args_1!(inputs);
58        unsafe {
59            let t = if self.modulo > input.shape()[self.axis] {
60                let mut shape: TVec<usize> = input.shape().into();
61                shape[self.axis] = 0;
62                Tensor::uninitialized_dt(input.datum_type(), &shape)?
63            } else {
64                let slice = ndarray::Slice::new(self.modulo as isize, None, self.stride);
65                unsafe fn do_slice<T: Datum>(
66                    t: &Tensor,
67                    axis: usize,
68                    slice: ndarray::Slice,
69                ) -> Tensor {
70                    let dt = t.datum_type();
71                    let mut t2 = t
72                        .to_array_view_unchecked::<T>()
73                        .slice_axis(Axis(axis), slice)
74                        .into_owned()
75                        .into_tensor();
76                    t2.set_datum_type(dt);
77                    t2
78                }
79                dispatch_datum_by_size!(do_slice(input.datum_type())(&*input, self.axis, slice))
80            };
81            Ok(tvec!(t.into_tvalue()))
82        }
83    }
84}
85
86impl TypedOp for Downsample {
87    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
88        ensure!(self.axis < inputs[0].rank());
89        ensure!(
90            self.modulo == 0 || self.stride > 0,
91            "non-zero modulo is only defined with forward strides"
92        );
93        let mut downed = inputs[0].without_value();
94        let down_len = self.transform_dim(&downed.shape[self.axis]);
95        downed.shape.set(self.axis, down_len);
96        Ok(tvec!(downed))
97    }
98
99    fn declutter(
100        &self,
101        model: &TypedModel,
102        node: &TypedNode,
103    ) -> TractResult<Option<TypedModelPatch>> {
104        if self.stride == 1 {
105            return Ok(Some(TypedModelPatch::replace_single_op(
106                model,
107                node,
108                &node.inputs,
109                Identity,
110            )?));
111        }
112        pull_downsample_up(model, node)
113            .with_context(|| format!("Pulling {} over {}", node, model.node(node.inputs[0].node)))
114    }
115
116    as_op!();
117}
118
119fn pull_downsample_up(
120    model: &TypedModel,
121    down_node: &TypedNode,
122) -> TractResult<Option<TypedModelPatch>> {
123    model.check_consistency()?;
124    let down_op = down_node.op_as::<Downsample>().unwrap();
125    if let Some(prec) = model.single_prec(down_node.id)? {
126        let (input_facts, output_facts) = model.node_facts(prec.id)?;
127        let axes_mapping = prec.op.axes_mapping(&input_facts, &output_facts)?;
128        debug!("Consider pull {:?} over {:?} (invariants: {:?})", down_op, prec, axes_mapping);
129        if let Some(slice_op) = prec.op_as::<ops::array::Slice>() {
130            if let Some(p) =
131                array::pull_downsample_over_slice(model, prec, slice_op, down_node, down_op)?
132            {
133                return Ok(Some(p));
134            }
135        } else if let Some(other_op) = prec.op_as::<AxisOp>() {
136            return array::pull_downsample_over_axis_op(model, prec, other_op, down_node, down_op);
137        } else if let Some(conv_op) = prec.op_as::<ops::cnn::conv::Conv>() {
138            return conv::fuse_downsample_into_conv(model, prec, conv_op, down_node, down_op);
139        } else if let Some(other_op) = prec.op_as::<ops::scan::Scan>() {
140            return scan::pull_downsample_over_scan(model, prec, other_op, down_node, down_op);
141        }
142        if prec.outputs.len() > 1 || prec.inputs.len() == 0 {
143            return Ok(None);
144        }
145        let axis_info = axes_mapping.axis((InOut::Out(0), down_op.axis))?;
146        let mut patch = TypedModelPatch::default();
147        let mut inputs = vec![];
148        for (ix, (outlet, axis_info)) in prec.inputs.iter().zip(&axis_info.inputs).enumerate() {
149            let mut wire = patch.tap_model(model, *outlet)?;
150            if let &[axis] = &**axis_info {
151                if !patch.outlet_fact(wire)?.shape[axis].is_one() {
152                    let mut op = down_op.clone();
153                    op.axis = axis;
154                    wire = patch.wire_node(
155                        format!("{}.{}-{}", down_node.name, prec.name, ix),
156                        op,
157                        &[wire],
158                    )?[0];
159                }
160            } else {
161                return Ok(None);
162            }
163            inputs.push(wire);
164        }
165        let other = patch.wire_node(&prec.name, prec.op.clone(), &inputs)?;
166        patch.shunt_outside(model, OutletId::new(down_node.id, 0), other[0])?;
167        return Ok(Some(patch));
168    }
169    Ok(None)
170}
tract_core/ops/downsample/mod.rs

tract_core/ops/downsample/
mod.rs
