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!(self.modulo == 0 || self.stride > 0, "non-zero modulo is only defined with forward strides");
90        let mut downed = inputs[0].clone();
91        let down_len = self.transform_dim(&downed.shape[self.axis]);
92        downed.shape.set(self.axis, down_len);
93        Ok(tvec!(downed))
94    }
95
96    fn declutter(
97        &self,
98        model: &TypedModel,
99        node: &TypedNode,
100    ) -> TractResult<Option<TypedModelPatch>> {
101        if self.stride == 1 {
102            return Ok(Some(TypedModelPatch::replace_single_op(
103                model,
104                node,
105                &node.inputs,
106                Identity,
107            )?));
108        }
109        pull_downsample_up(model, node)
110            .with_context(|| format!("Pulling {} over {}", node, model.node(node.inputs[0].node)))
111    }
112
113    as_op!();
114}
115
116fn pull_downsample_up(
117    model: &TypedModel,
118    down_node: &TypedNode,
119) -> TractResult<Option<TypedModelPatch>> {
120    model.check_consistency()?;
121    let down_op = down_node.op_as::<Downsample>().unwrap();
122    if let Some(prec) = model.single_prec(down_node.id)? {
123        let (input_facts, output_facts) = model.node_facts(prec.id)?;
124        let axes_mapping = prec.op.axes_mapping(&input_facts, &output_facts)?;
125        debug!("Consider pull {:?} over {:?} (invariants: {:?})", down_op, prec, axes_mapping);
126        if let Some(slice_op) = prec.op_as::<ops::array::Slice>() {
127            if let Some(p) =
128                array::pull_downsample_over_slice(model, prec, slice_op, down_node, down_op)?
129            {
130                return Ok(Some(p));
131            }
132        } else if let Some(other_op) = prec.op_as::<AxisOp>() {
133            return array::pull_downsample_over_axis_op(model, prec, other_op, down_node, down_op);
134        } else if let Some(conv_op) = prec.op_as::<ops::cnn::conv::Conv>() {
135            return conv::fuse_downsample_into_conv(model, prec, conv_op, down_node, down_op);
136        } else if let Some(other_op) = prec.op_as::<ops::scan::Scan>() {
137            return scan::pull_downsample_over_scan(model, prec, other_op, down_node, down_op);
138        }
139        if prec.outputs.len() > 1 || prec.inputs.len() == 0 {
140            return Ok(None);
141        }
142        let axis_info = axes_mapping.axis((InOut::Out(0), down_op.axis))?;
143        let mut patch = TypedModelPatch::default();
144        let mut inputs = vec![];
145        for (ix, (outlet, axis_info)) in prec.inputs.iter().zip(&axis_info.inputs).enumerate() {
146            let mut wire = patch.tap_model(model, *outlet)?;
147            if let &[axis] = &**axis_info {
148                if !patch.outlet_fact(wire)?.shape[axis].is_one() {
149                    let mut op = down_op.clone();
150                    op.axis = axis;
151                    wire = patch.wire_node(
152                        format!("{}.{}-{}", down_node.name, prec.name, ix),
153                        op,
154                        &[wire],
155                    )?[0];
156                }
157            } else {
158                return Ok(None);
159            }
160            inputs.push(wire);
161        }
162        let other = patch.wire_node(&prec.name, prec.op.clone(), &inputs)?;
163        patch.shunt_outside(model, OutletId::new(down_node.id, 0), other[0])?;
164        return Ok(Some(patch));
165    }
166    Ok(None)
167}
tract_core/ops/downsample/mod.rs

tract_core/ops/downsample/
mod.rs
