ruda-nn 0.21.56

Ruda neural network layers, activation modules and losses.
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

use crate::{PaddingConfig3d, padding::dilated_kernel_size};
use ruda_model::{
    config::Config,
    module::{Content, DisplaySettings, Module, ModuleDisplay},
    tensor::{Int, Tensor, backend::Backend,
        module::{max_pool3d_padded, max_pool3d_with_indices_padded}},
};

/// Configuration for maximum pooling over depth, height and width.
#[derive(Config, Debug)]
pub struct MaxPool3dConfig {
    /// Kernel extents in depth, height and width.
    pub kernel_size: [usize; 3],
    /// Strides, defaulting to the kernel extents.
    #[config(default = "kernel_size")]
    pub strides: [usize; 3],
    /// Symmetric, asymmetric, valid or dynamic same padding.
    #[config(default = "PaddingConfig3d::Valid")]
    pub padding: PaddingConfig3d,
    /// Spacing between sampled kernel elements.
    #[config(default = "[1, 1, 1]")]
    pub dilation: [usize; 3],
    /// Whether to include partially covered final windows.
    #[config(default = "false")]
    pub ceil_mode: bool,
}

/// Maximum pooling of native `[batch, channels, depth, height, width]` states.
#[derive(Module, Clone, Debug)]
#[module(custom_display)]
pub struct MaxPool3d {
    /// Strides in depth, height and width.
    pub stride: [usize; 3],
    /// Kernel extents.
    pub kernel_size: [usize; 3],
    /// Padding policy.
    pub padding: PaddingConfig3d,
    /// Kernel dilation.
    pub dilation: [usize; 3],
    /// Whether to include partially covered final windows.
    pub ceil_mode: bool,
}

impl MaxPool3dConfig {
    /// Construct the pooling layer without allocating model parameters.
    pub fn init(&self) -> MaxPool3d {
        MaxPool3d {
            stride: self.strides,
            kernel_size: self.kernel_size,
            padding: self.padding.clone(),
            dilation: self.dilation,
            ceil_mode: self.ceil_mode,
        }
    }
}

impl MaxPool3d {
    /// Pool activations and return flattened positions in the unpadded input volume.
    pub fn forward_with_indices<B: Backend>(&self, input: Tensor<B, 5>)
        -> (Tensor<B, 5>, Tensor<B, 5, Int>) {
        let [_, _, depth, height, width] = input.dims();
        let effective = core::array::from_fn(|axis| {
            dilated_kernel_size(self.kernel_size[axis], self.dilation[axis])
        });
        let padding = self.padding.calculate_padding_3d_pairs(
            &[depth, height, width], &effective, &self.stride,
        );
        max_pool3d_with_indices_padded(input, self.kernel_size, self.stride, padding,
            self.dilation, self.ceil_mode)
    }

    /// Reduce a volume using the configured native pooling operations.
    pub fn forward<B: Backend>(&self, input: Tensor<B, 5>) -> Tensor<B, 5> {
        let [_, _, depth, height, width] = input.dims();
        let effective = core::array::from_fn(|axis| {
            dilated_kernel_size(self.kernel_size[axis], self.dilation[axis])
        });
        let pairs = self.padding.calculate_padding_3d_pairs(
            &[depth, height, width], &effective, &self.stride,
        );
        max_pool3d_padded(input, self.kernel_size, self.stride, pairs, self.dilation, self.ceil_mode)
    }
}

impl ModuleDisplay for MaxPool3d {
    fn custom_settings(&self) -> Option<DisplaySettings> {
        DisplaySettings::new().with_new_line_after_attribute(false).optional()
    }

    fn custom_content(&self, content: Content) -> Option<Content> {
        content.add_debug_attribute("kernel_size", &self.kernel_size)
            .add_debug_attribute("stride", &self.stride)
            .add_debug_attribute("padding", &self.padding)
            .add_debug_attribute("dilation", &self.dilation)
            .add("ceil_mode", &self.ceil_mode).optional()
    }
}