hodu_nn/modules/
pooling.rs

1use crate::compat::*;
2use crate::module::Module;
3use hodu_core::{error::HoduResult, tensor::Tensor};
4
5#[derive(Module, Clone)]
6pub struct AdaptiveAvgPool1D {
7    output_size: u32,
8}
9
10impl AdaptiveAvgPool1D {
11    pub fn new(output_size: u32) -> Self {
12        Self { output_size }
13    }
14
15    pub fn forward(&self, input: &Tensor) -> HoduResult<Tensor> {
16        // AdaptiveAvgPool1D: input [N, C, L]
17        let input_shape = input.shape();
18        let rank = input_shape.ndim();
19
20        if rank != 3 {
21            return Err(hodu_core::error::HoduError::InternalError(format!(
22                "AdaptiveAvgPool1D expects 3D input [N, C, L], got {}D",
23                rank
24            )));
25        }
26
27        let input_length = input_shape[2];
28
29        let stride = (input_length as f64 / self.output_size as f64).floor() as u32;
30        let kernel_size = input_length - (self.output_size - 1) * stride;
31
32        let padding = vec![(0, 0), (0, 0), (0, 0)];
33        let window_shape = vec![1, 1, kernel_size];
34        let strides = vec![1, 1, stride];
35
36        input.reduce_window(&window_shape, &strides, &padding, "mean")
37    }
38
39    pub fn parameters(&mut self) -> Vec<&mut Tensor> {
40        vec![]
41    }
42}
43
44#[derive(Module, Clone)]
45pub struct AdaptiveAvgPool2D {
46    output_size: (u32, u32),
47}
48
49impl AdaptiveAvgPool2D {
50    pub fn new(output_size: (u32, u32)) -> Self {
51        Self { output_size }
52    }
53
54    pub fn forward(&self, input: &Tensor) -> HoduResult<Tensor> {
55        // AdaptiveAvgPool2D: input [N, C, H, W]
56        let input_shape = input.shape();
57        let rank = input_shape.ndim();
58
59        if rank != 4 {
60            return Err(hodu_core::error::HoduError::InternalError(format!(
61                "AdaptiveAvgPool2D expects 4D input [N, C, H, W], got {}D",
62                rank
63            )));
64        }
65
66        let input_height = input_shape[2];
67        let input_width = input_shape[3];
68
69        let stride_h = (input_height as f64 / self.output_size.0 as f64).floor() as u32;
70        let stride_w = (input_width as f64 / self.output_size.1 as f64).floor() as u32;
71        let kernel_h = input_height - (self.output_size.0 - 1) * stride_h;
72        let kernel_w = input_width - (self.output_size.1 - 1) * stride_w;
73
74        let padding = vec![(0, 0), (0, 0), (0, 0), (0, 0)];
75        let window_shape = vec![1, 1, kernel_h, kernel_w];
76        let strides = vec![1, 1, stride_h, stride_w];
77
78        input.reduce_window(&window_shape, &strides, &padding, "mean")
79    }
80
81    pub fn parameters(&mut self) -> Vec<&mut Tensor> {
82        vec![]
83    }
84}
85
86#[derive(Module, Clone)]
87pub struct AdaptiveAvgPool3D {
88    output_size: (u32, u32, u32),
89}
90
91impl AdaptiveAvgPool3D {
92    pub fn new(output_size: (u32, u32, u32)) -> Self {
93        Self { output_size }
94    }
95
96    pub fn forward(&self, input: &Tensor) -> HoduResult<Tensor> {
97        // AdaptiveAvgPool3D: input [N, C, D, H, W]
98        let input_shape = input.shape();
99        let rank = input_shape.ndim();
100
101        if rank != 5 {
102            return Err(hodu_core::error::HoduError::InternalError(format!(
103                "AdaptiveAvgPool3D expects 5D input [N, C, D, H, W], got {}D",
104                rank
105            )));
106        }
107
108        let input_depth = input_shape[2];
109        let input_height = input_shape[3];
110        let input_width = input_shape[4];
111
112        let stride_d = (input_depth as f64 / self.output_size.0 as f64).floor() as u32;
113        let stride_h = (input_height as f64 / self.output_size.1 as f64).floor() as u32;
114        let stride_w = (input_width as f64 / self.output_size.2 as f64).floor() as u32;
115        let kernel_d = input_depth - (self.output_size.0 - 1) * stride_d;
116        let kernel_h = input_height - (self.output_size.1 - 1) * stride_h;
117        let kernel_w = input_width - (self.output_size.2 - 1) * stride_w;
118
119        let padding = vec![(0, 0), (0, 0), (0, 0), (0, 0), (0, 0)];
120        let window_shape = vec![1, 1, kernel_d, kernel_h, kernel_w];
121        let strides = vec![1, 1, stride_d, stride_h, stride_w];
122
123        input.reduce_window(&window_shape, &strides, &padding, "mean")
124    }
125
126    pub fn parameters(&mut self) -> Vec<&mut Tensor> {
127        vec![]
128    }
129}
130
131#[derive(Module, Clone)]
132pub struct AdaptiveMaxPool1D {
133    output_size: u32,
134}
135
136impl AdaptiveMaxPool1D {
137    pub fn new(output_size: u32) -> Self {
138        Self { output_size }
139    }
140
141    pub fn forward(&self, input: &Tensor) -> HoduResult<Tensor> {
142        // AdaptiveMaxPool1D: input [N, C, L]
143        let input_shape = input.shape();
144        let rank = input_shape.ndim();
145
146        if rank != 3 {
147            return Err(hodu_core::error::HoduError::InternalError(format!(
148                "AdaptiveMaxPool1D expects 3D input [N, C, L], got {}D",
149                rank
150            )));
151        }
152
153        let input_length = input_shape[2];
154
155        // Calculate kernel_size and stride for adaptive pooling
156        let stride = (input_length as f64 / self.output_size as f64).floor() as u32;
157        let kernel_size = input_length - (self.output_size - 1) * stride;
158
159        let padding = vec![(0, 0), (0, 0), (0, 0)];
160        let window_shape = vec![1, 1, kernel_size];
161        let strides = vec![1, 1, stride];
162
163        input.reduce_window(&window_shape, &strides, &padding, "max")
164    }
165
166    pub fn parameters(&mut self) -> Vec<&mut Tensor> {
167        vec![]
168    }
169}
170
171#[derive(Module, Clone)]
172pub struct AdaptiveMaxPool2D {
173    output_size: (u32, u32),
174}
175
176impl AdaptiveMaxPool2D {
177    pub fn new(output_size: (u32, u32)) -> Self {
178        Self { output_size }
179    }
180
181    pub fn forward(&self, input: &Tensor) -> HoduResult<Tensor> {
182        // AdaptiveMaxPool2D: input [N, C, H, W]
183        let input_shape = input.shape();
184        let rank = input_shape.ndim();
185
186        if rank != 4 {
187            return Err(hodu_core::error::HoduError::InternalError(format!(
188                "AdaptiveMaxPool2D expects 4D input [N, C, H, W], got {}D",
189                rank
190            )));
191        }
192
193        let input_height = input_shape[2];
194        let input_width = input_shape[3];
195
196        let stride_h = (input_height as f64 / self.output_size.0 as f64).floor() as u32;
197        let stride_w = (input_width as f64 / self.output_size.1 as f64).floor() as u32;
198        let kernel_h = input_height - (self.output_size.0 - 1) * stride_h;
199        let kernel_w = input_width - (self.output_size.1 - 1) * stride_w;
200
201        let padding = vec![(0, 0), (0, 0), (0, 0), (0, 0)];
202        let window_shape = vec![1, 1, kernel_h, kernel_w];
203        let strides = vec![1, 1, stride_h, stride_w];
204
205        input.reduce_window(&window_shape, &strides, &padding, "max")
206    }
207
208    pub fn parameters(&mut self) -> Vec<&mut Tensor> {
209        vec![]
210    }
211}
212
213#[derive(Module, Clone)]
214pub struct AdaptiveMaxPool3D {
215    output_size: (u32, u32, u32),
216}
217
218impl AdaptiveMaxPool3D {
219    pub fn new(output_size: (u32, u32, u32)) -> Self {
220        Self { output_size }
221    }
222
223    pub fn forward(&self, input: &Tensor) -> HoduResult<Tensor> {
224        // AdaptiveMaxPool3D: input [N, C, D, H, W]
225        let input_shape = input.shape();
226        let rank = input_shape.ndim();
227
228        if rank != 5 {
229            return Err(hodu_core::error::HoduError::InternalError(format!(
230                "AdaptiveMaxPool3D expects 5D input [N, C, D, H, W], got {}D",
231                rank
232            )));
233        }
234
235        let input_depth = input_shape[2];
236        let input_height = input_shape[3];
237        let input_width = input_shape[4];
238
239        let stride_d = (input_depth as f64 / self.output_size.0 as f64).floor() as u32;
240        let stride_h = (input_height as f64 / self.output_size.1 as f64).floor() as u32;
241        let stride_w = (input_width as f64 / self.output_size.2 as f64).floor() as u32;
242        let kernel_d = input_depth - (self.output_size.0 - 1) * stride_d;
243        let kernel_h = input_height - (self.output_size.1 - 1) * stride_h;
244        let kernel_w = input_width - (self.output_size.2 - 1) * stride_w;
245
246        let padding = vec![(0, 0), (0, 0), (0, 0), (0, 0), (0, 0)];
247        let window_shape = vec![1, 1, kernel_d, kernel_h, kernel_w];
248        let strides = vec![1, 1, stride_d, stride_h, stride_w];
249
250        input.reduce_window(&window_shape, &strides, &padding, "max")
251    }
252
253    pub fn parameters(&mut self) -> Vec<&mut Tensor> {
254        vec![]
255    }
256}
257
258#[derive(Module, Clone)]
259pub struct AvgPool1D {
260    kernel_size: u32,
261    stride: u32,
262    padding: u32,
263}
264
265impl AvgPool1D {
266    pub fn new(kernel_size: u32, stride: u32, padding: u32) -> Self {
267        Self {
268            kernel_size,
269            stride,
270            padding,
271        }
272    }
273
274    pub fn forward(&self, input: &Tensor) -> HoduResult<Tensor> {
275        // AvgPool1D: input [N, C, L]
276        let input_shape = input.shape();
277        let rank = input_shape.ndim();
278
279        if rank != 3 {
280            return Err(hodu_core::error::HoduError::InternalError(format!(
281                "AvgPool1D expects 3D input [N, C, L], got {}D",
282                rank
283            )));
284        }
285
286        let padding = vec![(0, 0), (0, 0), (self.padding, self.padding)];
287        let window_shape = vec![1, 1, self.kernel_size];
288        let strides = vec![1, 1, self.stride];
289
290        input.reduce_window(&window_shape, &strides, &padding, "mean")
291    }
292
293    pub fn parameters(&mut self) -> Vec<&mut Tensor> {
294        vec![]
295    }
296}
297
298#[derive(Module, Clone)]
299pub struct AvgPool2D {
300    kernel_size: u32,
301    stride: u32,
302    padding: u32,
303}
304
305impl AvgPool2D {
306    pub fn new(kernel_size: u32, stride: u32, padding: u32) -> Self {
307        Self {
308            kernel_size,
309            stride,
310            padding,
311        }
312    }
313
314    pub fn forward(&self, input: &Tensor) -> HoduResult<Tensor> {
315        // AvgPool2D: input [N, C, H, W]
316        let input_shape = input.shape();
317        let rank = input_shape.ndim();
318
319        if rank != 4 {
320            return Err(hodu_core::error::HoduError::InternalError(format!(
321                "AvgPool2D expects 4D input [N, C, H, W], got {}D",
322                rank
323            )));
324        }
325
326        let padding = vec![
327            (0, 0),
328            (0, 0),
329            (self.padding, self.padding),
330            (self.padding, self.padding),
331        ];
332        let window_shape = vec![1, 1, self.kernel_size, self.kernel_size];
333        let strides = vec![1, 1, self.stride, self.stride];
334
335        input.reduce_window(&window_shape, &strides, &padding, "mean")
336    }
337
338    pub fn parameters(&mut self) -> Vec<&mut Tensor> {
339        vec![]
340    }
341}
342
343#[derive(Module, Clone)]
344pub struct AvgPool3D {
345    kernel_size: u32,
346    stride: u32,
347    padding: u32,
348}
349
350impl AvgPool3D {
351    pub fn new(kernel_size: u32, stride: u32, padding: u32) -> Self {
352        Self {
353            kernel_size,
354            stride,
355            padding,
356        }
357    }
358
359    pub fn forward(&self, input: &Tensor) -> HoduResult<Tensor> {
360        // AvgPool3D: input [N, C, D, H, W]
361        let input_shape = input.shape();
362        let rank = input_shape.ndim();
363
364        if rank != 5 {
365            return Err(hodu_core::error::HoduError::InternalError(format!(
366                "AvgPool3D expects 5D input [N, C, D, H, W], got {}D",
367                rank
368            )));
369        }
370
371        let padding = vec![
372            (0, 0),
373            (0, 0),
374            (self.padding, self.padding),
375            (self.padding, self.padding),
376            (self.padding, self.padding),
377        ];
378        let window_shape = vec![1, 1, self.kernel_size, self.kernel_size, self.kernel_size];
379        let strides = vec![1, 1, self.stride, self.stride, self.stride];
380
381        input.reduce_window(&window_shape, &strides, &padding, "mean")
382    }
383
384    pub fn parameters(&mut self) -> Vec<&mut Tensor> {
385        vec![]
386    }
387}
388
389#[derive(Module, Clone)]
390pub struct MaxPool1D {
391    kernel_size: u32,
392    stride: u32,
393    padding: u32,
394}
395
396impl MaxPool1D {
397    pub fn new(kernel_size: u32, stride: u32, padding: u32) -> Self {
398        Self {
399            kernel_size,
400            stride,
401            padding,
402        }
403    }
404
405    pub fn forward(&self, input: &Tensor) -> HoduResult<Tensor> {
406        // MaxPool1D: input [N, C, L]
407        let input_shape = input.shape();
408        let rank = input_shape.ndim();
409
410        if rank != 3 {
411            return Err(hodu_core::error::HoduError::InternalError(format!(
412                "MaxPool1D expects 3D input [N, C, L], got {}D",
413                rank
414            )));
415        }
416
417        // Build padding, window_shape, and strides for all dimensions
418        let padding = vec![(0, 0), (0, 0), (self.padding, self.padding)];
419        let window_shape = vec![1, 1, self.kernel_size];
420        let strides = vec![1, 1, self.stride];
421
422        input.reduce_window(&window_shape, &strides, &padding, "max")
423    }
424
425    pub fn parameters(&mut self) -> Vec<&mut Tensor> {
426        vec![]
427    }
428}
429
430#[derive(Module, Clone)]
431pub struct MaxPool2D {
432    kernel_size: u32,
433    stride: u32,
434    padding: u32,
435}
436
437impl MaxPool2D {
438    pub fn new(kernel_size: u32, stride: u32, padding: u32) -> Self {
439        Self {
440            kernel_size,
441            stride,
442            padding,
443        }
444    }
445
446    pub fn forward(&self, input: &Tensor) -> HoduResult<Tensor> {
447        // MaxPool2D: input [N, C, H, W]
448        let input_shape = input.shape();
449        let rank = input_shape.ndim();
450
451        if rank != 4 {
452            return Err(hodu_core::error::HoduError::InternalError(format!(
453                "MaxPool2D expects 4D input [N, C, H, W], got {}D",
454                rank
455            )));
456        }
457
458        let padding = vec![
459            (0, 0),
460            (0, 0),
461            (self.padding, self.padding),
462            (self.padding, self.padding),
463        ];
464        let window_shape = vec![1, 1, self.kernel_size, self.kernel_size];
465        let strides = vec![1, 1, self.stride, self.stride];
466
467        input.reduce_window(&window_shape, &strides, &padding, "max")
468    }
469
470    pub fn parameters(&mut self) -> Vec<&mut Tensor> {
471        vec![]
472    }
473}
474
475#[derive(Module, Clone)]
476pub struct MaxPool3D {
477    kernel_size: u32,
478    stride: u32,
479    padding: u32,
480}
481
482impl MaxPool3D {
483    pub fn new(kernel_size: u32, stride: u32, padding: u32) -> Self {
484        Self {
485            kernel_size,
486            stride,
487            padding,
488        }
489    }
490
491    pub fn forward(&self, input: &Tensor) -> HoduResult<Tensor> {
492        // MaxPool3D: input [N, C, D, H, W]
493        let input_shape = input.shape();
494        let rank = input_shape.ndim();
495
496        if rank != 5 {
497            return Err(hodu_core::error::HoduError::InternalError(format!(
498                "MaxPool3D expects 5D input [N, C, D, H, W], got {}D",
499                rank
500            )));
501        }
502
503        let padding = vec![
504            (0, 0),
505            (0, 0),
506            (self.padding, self.padding),
507            (self.padding, self.padding),
508            (self.padding, self.padding),
509        ];
510        let window_shape = vec![1, 1, self.kernel_size, self.kernel_size, self.kernel_size];
511        let strides = vec![1, 1, self.stride, self.stride, self.stride];
512
513        input.reduce_window(&window_shape, &strides, &padding, "max")
514    }
515
516    pub fn parameters(&mut self) -> Vec<&mut Tensor> {
517        vec![]
518    }
519}
hodu_nn/modules/pooling.rs

hodu_nn/modules/
pooling.rs
