rustyml 0.17.0

Machine learning models, neural networks, metrics, and data utilities in Rust
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150

//! 3D upsampling layer that enlarges a volume by a whole-number factor per axis
//!
//! [`UpSampling3D`] enlarges the depth axis, the height axis, and the width axis of a
//! `[batch, dim1, dim2, dim3, channels]` tensor. It holds no parameter and takes no interpolation
//! mode. Each output position copies its source position through `resize_engine`, which also
//! runs the matching backward pass.

use crate::error::Error;
use crate::neural_network::layers::ParamCounts;
use crate::neural_network::layers::upsampling::resize_engine::{
    upsample_backward, upsample_forward, upsample_output_shape, validate_factors,
};
use crate::neural_network::layers::upsampling::{Factor3D, Interpolation};
use crate::neural_network::layers::validation::start_build;
use crate::neural_network::layers::{
    built_layer_shape_functions, no_trainable_parameters_layer_functions,
};
use crate::neural_network::traits::{LayerBase, UnaryLayer};
use crate::neural_network::{Ctx, Shape, Tensor};

/// Enlarges the 3 spatial axes of a rank-5 tensor
///
/// The input shape is `[batch_size, dim1, dim2, dim3, channels]`. The output shape multiplies
/// each spatial extent by its factor. The batch axis and the channel axis pass through
/// unchanged
///
/// The layer holds no parameter. Each output position copies the input position it came from,
/// so the layer takes no interpolation argument
///
/// The layer is the decoder counterpart of
/// [`MaxPooling3D`](crate::neural_network::layers::pooling::MaxPooling3D) and
/// [`AveragePooling3D`](crate::neural_network::layers::pooling::AveragePooling3D). A volume model
/// pairs each pooling stage with 1 upsampling stage of the same factor. Examples include a model
/// over a medical scan or a video clip
///
/// # Examples
///
/// ```rust
/// use rustyml::neural_network::Shape;
/// use rustyml::neural_network::sequential::SequentialBuilder;
/// use rustyml::neural_network::layers::*;
/// use rustyml::neural_network::optimizers::*;
/// use rustyml::neural_network::losses::*;
/// use ndarray::Array5;
///
/// // A rank-5 input: 1 sample, a 2x1x1 volume, 2 channels.
/// let x = Array5::from_shape_vec((1, 2, 1, 1, 2), vec![1.0, 2.0, 3.0, 4.0])
///     .unwrap()
///     .into_dyn();
///
/// let mut model = SequentialBuilder::new()
///     .add(UpSampling3D::new((2, 3, 1)).unwrap())
///     .build(&Shape::known(x.shape()))
///     .unwrap();
/// model.compile(SGD::new(0.01, 0.0, false, 0.0).unwrap(), MeanSquaredError::new());
///
/// // View model structure.
/// model.summary();
///
/// let larger = model.predict(&x).unwrap();
///
/// // Each spatial axis grows by its own factor, and the channel axis stays put.
/// assert_eq!(larger.shape(), &[1, 4, 3, 1, 2]);
/// assert_eq!(larger[[0, 0, 2, 0, 1]], 2.0);
/// assert_eq!(larger[[0, 3, 0, 0, 0]], 3.0);
/// ```
#[derive(Debug)]
pub struct UpSampling3D {
    /// Factor each spatial axis grows by
    size: Factor3D,
    /// Shape the layer was built for, batch axis first. `None` before the build
    built: Option<Shape>,
}

impl UpSampling3D {
    /// Creates a new UpSampling3D layer
    ///
    /// # Parameters
    ///
    /// - `size` - Factor each spatial axis grows by. An integer gives the same factor to all 3
    ///   axes. A triple names the factor of each axis. See [`Factor3D`]
    ///
    /// # Returns
    ///
    /// - `Result<Self, Error>` - New `UpSampling3D` layer instance
    ///
    /// # Errors
    ///
    /// - `Error::InvalidParameter` - If any factor is 0
    pub fn new(size: impl Into<Factor3D>) -> Result<Self, Error> {
        let size = size.into();
        validate_factors(&size.0)?;
        Ok(UpSampling3D { size, built: None })
    }
}

impl LayerBase for UpSampling3D {
    fn layer_type(&self) -> &str {
        "UpSampling3D"
    }

    built_layer_shape_functions!();

    no_trainable_parameters_layer_functions!();
}

impl UnaryLayer for UpSampling3D {
    /// Records the shape the layer enlarges. The layer holds no array, so nothing is allocated.
    /// The shape algebra checks the rank
    fn build(&mut self, input: &Shape) -> Result<(), Error> {
        let Some(built) = start_build(&self.built, "UpSampling3D", input)? else {
            return Ok(());
        };
        self.compute_output_shape(&built)?;
        self.built = Some(built);
        Ok(())
    }

    fn forward(&self, input: &Tensor, ctx: &mut Ctx) -> Result<Tensor, Error> {
        let output = upsample_forward(
            input,
            &self.size.0,
            Interpolation::Nearest,
            5,
            "UpSampling3D",
        )?;

        if ctx.is_training() {
            ctx.push_cache("UpSampling3D", input.shape().to_vec());
        }

        Ok(output)
    }

    fn backward(&self, grad_output: &Tensor, ctx: &mut Ctx) -> Result<Tensor, Error> {
        let input_shape: Vec<usize> = ctx.pop_cache("UpSampling3D")?;

        upsample_backward(
            grad_output,
            Some(input_shape.as_slice()),
            &self.size.0,
            Interpolation::Nearest,
            "UpSampling3D",
        )
    }

    fn compute_output_shape(&self, input: &Shape) -> Result<Shape, Error> {
        upsample_output_shape(input, &self.size.0, "UpSampling3D")
    }
}