xref: /aosp_15_r20/external/pytorch/torch/csrc/api/include/torch/nn/options/pooling.h (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1 #pragma once
2 
3 #include <torch/arg.h>
4 #include <torch/csrc/Export.h>
5 #include <torch/expanding_array.h>
6 #include <torch/types.h>
7 
8 namespace torch {
9 namespace nn {
10 
11 /// Options for a `D`-dimensional avgpool module.
12 template <size_t D>
13 struct AvgPoolOptions {
AvgPoolOptionsAvgPoolOptions14   AvgPoolOptions(ExpandingArray<D> kernel_size)
15       : kernel_size_(kernel_size), stride_(kernel_size) {}
16 
17   /// the size of the window to take an average over
18   TORCH_ARG(ExpandingArray<D>, kernel_size);
19 
20   /// the stride of the window. Default value is `kernel_size`
21   TORCH_ARG(ExpandingArray<D>, stride);
22 
23   /// implicit zero padding to be added on both sides
24   TORCH_ARG(ExpandingArray<D>, padding) = 0;
25 
26   /// when True, will use `ceil` instead of `floor` to compute the output shape
27   TORCH_ARG(bool, ceil_mode) = false;
28 
29   /// when True, will include the zero-padding in the averaging calculation
30   TORCH_ARG(bool, count_include_pad) = true;
31 
32   /// if specified, it will be used as divisor, otherwise size of the pooling
33   /// region will be used.
34 
35   TORCH_ARG(std::optional<int64_t>, divisor_override) = std::nullopt;
36 };
37 
38 /// `AvgPoolOptions` specialized for the `AvgPool1d` module.
39 ///
40 /// Example:
41 /// ```
42 /// AvgPool1d model(AvgPool1dOptions(3).stride(2));
43 /// ```
44 using AvgPool1dOptions = AvgPoolOptions<1>;
45 
46 /// `AvgPoolOptions` specialized for the `AvgPool2d` module.
47 ///
48 /// Example:
49 /// ```
50 /// AvgPool2d model(AvgPool2dOptions({3, 2}).stride({2, 2}));
51 /// ```
52 using AvgPool2dOptions = AvgPoolOptions<2>;
53 
54 /// `AvgPoolOptions` specialized for the `AvgPool3d` module.
55 ///
56 /// Example:
57 /// ```
58 /// AvgPool3d model(AvgPool3dOptions(5).stride(2));
59 /// ```
60 using AvgPool3dOptions = AvgPoolOptions<3>;
61 
62 namespace functional {
63 /// Options for `torch::nn::functional::avg_pool1d`.
64 ///
65 /// See the documentation for `torch::nn::AvgPool1dOptions` class to learn what
66 /// arguments are supported.
67 ///
68 /// Example:
69 /// ```
70 /// namespace F = torch::nn::functional;
71 /// F::avg_pool1d(x, F::AvgPool1dFuncOptions(3).stride(2));
72 /// ```
73 using AvgPool1dFuncOptions = AvgPool1dOptions;
74 } // namespace functional
75 
76 namespace functional {
77 /// Options for `torch::nn::functional::avg_pool2d`.
78 ///
79 /// See the documentation for `torch::nn::AvgPool2dOptions` class to learn what
80 /// arguments are supported.
81 ///
82 /// Example:
83 /// ```
84 /// namespace F = torch::nn::functional;
85 /// F::avg_pool2d(x, F::AvgPool2dFuncOptions(3).stride(2));
86 /// ```
87 using AvgPool2dFuncOptions = AvgPool2dOptions;
88 } // namespace functional
89 
90 namespace functional {
91 /// Options for `torch::nn::functional::avg_pool3d`.
92 ///
93 /// See the documentation for `torch::nn::AvgPool3dOptions` class to learn what
94 /// arguments are supported.
95 ///
96 /// Example:
97 /// ```
98 /// namespace F = torch::nn::functional;
99 /// F::avg_pool3d(x, F::AvgPool3dFuncOptions(3).stride(2));
100 /// ```
101 using AvgPool3dFuncOptions = AvgPool3dOptions;
102 } // namespace functional
103 
104 // ============================================================================
105 
106 /// Options for a `D`-dimensional maxpool module.
107 template <size_t D>
108 struct MaxPoolOptions {
MaxPoolOptionsMaxPoolOptions109   MaxPoolOptions(ExpandingArray<D> kernel_size)
110       : kernel_size_(kernel_size), stride_(kernel_size) {}
111 
112   /// the size of the window to take a max over
113   TORCH_ARG(ExpandingArray<D>, kernel_size);
114 
115   /// the stride of the window. Default value is `kernel_size
116   TORCH_ARG(ExpandingArray<D>, stride);
117 
118   /// implicit zero padding to be added on both sides
119   TORCH_ARG(ExpandingArray<D>, padding) = 0;
120 
121   /// a parameter that controls the stride of elements in the window
122   TORCH_ARG(ExpandingArray<D>, dilation) = 1;
123 
124   /// when True, will use `ceil` instead of `floor` to compute the output shape
125   TORCH_ARG(bool, ceil_mode) = false;
126 };
127 
128 /// `MaxPoolOptions` specialized for the `MaxPool1d` module.
129 ///
130 /// Example:
131 /// ```
132 /// MaxPool1d model(MaxPool1dOptions(3).stride(2));
133 /// ```
134 using MaxPool1dOptions = MaxPoolOptions<1>;
135 
136 /// `MaxPoolOptions` specialized for the `MaxPool2d` module.
137 ///
138 /// Example:
139 /// ```
140 /// MaxPool2d model(MaxPool2dOptions({3, 2}).stride({2, 2}));
141 /// ```
142 using MaxPool2dOptions = MaxPoolOptions<2>;
143 
144 /// `MaxPoolOptions` specialized for the `MaxPool3d` module.
145 ///
146 /// Example:
147 /// ```
148 /// MaxPool3d model(MaxPool3dOptions(3).stride(2));
149 /// ```
150 using MaxPool3dOptions = MaxPoolOptions<3>;
151 
152 namespace functional {
153 /// Options for `torch::nn::functional::max_pool1d` and
154 /// `torch::nn::functional::max_pool1d_with_indices`.
155 ///
156 /// Example:
157 /// ```
158 /// namespace F = torch::nn::functional;
159 /// F::max_pool1d(x, F::MaxPool1dFuncOptions(3).stride(2));
160 /// ```
161 using MaxPool1dFuncOptions = MaxPool1dOptions;
162 } // namespace functional
163 
164 namespace functional {
165 /// Options for `torch::nn::functional::max_pool2d` and
166 /// `torch::nn::functional::max_pool2d_with_indices`.
167 ///
168 /// Example:
169 /// ```
170 /// namespace F = torch::nn::functional;
171 /// F::max_pool2d(x, F::MaxPool2dFuncOptions(3).stride(2));
172 /// ```
173 using MaxPool2dFuncOptions = MaxPool2dOptions;
174 } // namespace functional
175 
176 namespace functional {
177 /// Options for `torch::nn::functional::max_pool3d` and
178 /// `torch::nn::functional::max_pool3d_with_indices`.
179 ///
180 /// Example:
181 /// ```
182 /// namespace F = torch::nn::functional;
183 /// F::max_pool3d(x, F::MaxPool3dFuncOptions(3).stride(2));
184 /// ```
185 using MaxPool3dFuncOptions = MaxPool3dOptions;
186 } // namespace functional
187 
188 // ============================================================================
189 
190 /// Options for a `D`-dimensional adaptive maxpool module.
191 template <typename output_size_t>
192 struct AdaptiveMaxPoolOptions {
AdaptiveMaxPoolOptionsAdaptiveMaxPoolOptions193   AdaptiveMaxPoolOptions(output_size_t output_size)
194       : output_size_(output_size) {}
195 
196   /// the target output size
197   TORCH_ARG(output_size_t, output_size);
198 };
199 
200 /// `AdaptiveMaxPoolOptions` specialized for the `AdaptiveMaxPool1d` module.
201 ///
202 /// Example:
203 /// ```
204 /// AdaptiveMaxPool1d model(AdaptiveMaxPool1dOptions(3));
205 /// ```
206 using AdaptiveMaxPool1dOptions = AdaptiveMaxPoolOptions<ExpandingArray<1>>;
207 
208 /// `AdaptiveMaxPoolOptions` specialized for the `AdaptiveMaxPool2d` module.
209 ///
210 /// Example:
211 /// ```
212 /// AdaptiveMaxPool2d model(AdaptiveMaxPool2dOptions({3, 2}));
213 /// ```
214 using AdaptiveMaxPool2dOptions =
215     AdaptiveMaxPoolOptions<ExpandingArrayWithOptionalElem<2>>;
216 
217 /// `AdaptiveMaxPoolOptions` specialized for the `AdaptiveMaxPool3d` module.
218 ///
219 /// Example:
220 /// ```
221 /// AdaptiveMaxPool3d model(AdaptiveMaxPool3dOptions(3));
222 /// ```
223 using AdaptiveMaxPool3dOptions =
224     AdaptiveMaxPoolOptions<ExpandingArrayWithOptionalElem<3>>;
225 
226 namespace functional {
227 /// Options for `torch::nn::functional::adaptive_max_pool1d` and
228 /// `torch::nn::functional::adaptive_max_pool1d_with_indices`
229 ///
230 /// Example:
231 /// ```
232 /// namespace F = torch::nn::functional;
233 /// F::adaptive_max_pool1d(x, F::AdaptiveMaxPool1dFuncOptions(3));
234 /// ```
235 using AdaptiveMaxPool1dFuncOptions = AdaptiveMaxPool1dOptions;
236 } // namespace functional
237 
238 namespace functional {
239 /// Options for `torch::nn::functional::adaptive_max_pool2d` and
240 /// `torch::nn::functional::adaptive_max_pool2d_with_indices`
241 ///
242 /// Example:
243 /// ```
244 /// namespace F = torch::nn::functional;
245 /// F::adaptive_max_pool2d(x, F::AdaptiveMaxPool2dFuncOptions(3));
246 /// ```
247 using AdaptiveMaxPool2dFuncOptions = AdaptiveMaxPool2dOptions;
248 } // namespace functional
249 
250 namespace functional {
251 /// Options for `torch::nn::functional::adaptive_max_pool3d` and
252 /// `torch::nn::functional::adaptive_max_pool3d_with_indices`
253 ///
254 /// Example:
255 /// ```
256 /// namespace F = torch::nn::functional;
257 /// F::adaptive_max_pool3d(x, F::AdaptiveMaxPool3dFuncOptions(3));
258 /// ```
259 using AdaptiveMaxPool3dFuncOptions = AdaptiveMaxPool3dOptions;
260 } // namespace functional
261 
262 // ============================================================================
263 
264 /// Options for a `D`-dimensional adaptive avgpool module.
265 template <typename output_size_t>
266 struct AdaptiveAvgPoolOptions {
AdaptiveAvgPoolOptionsAdaptiveAvgPoolOptions267   AdaptiveAvgPoolOptions(output_size_t output_size)
268       : output_size_(output_size) {}
269 
270   /// the target output size
271   TORCH_ARG(output_size_t, output_size);
272 };
273 
274 /// `AdaptiveAvgPoolOptions` specialized for the `AdaptiveAvgPool1d` module.
275 ///
276 /// Example:
277 /// ```
278 /// AdaptiveAvgPool1d model(AdaptiveAvgPool1dOptions(5));
279 /// ```
280 using AdaptiveAvgPool1dOptions = AdaptiveAvgPoolOptions<ExpandingArray<1>>;
281 
282 /// `AdaptiveAvgPoolOptions` specialized for the `AdaptiveAvgPool2d` module.
283 ///
284 /// Example:
285 /// ```
286 /// AdaptiveAvgPool2d model(AdaptiveAvgPool2dOptions({3, 2}));
287 /// ```
288 using AdaptiveAvgPool2dOptions =
289     AdaptiveAvgPoolOptions<ExpandingArrayWithOptionalElem<2>>;
290 
291 /// `AdaptiveAvgPoolOptions` specialized for the `AdaptiveAvgPool3d` module.
292 ///
293 /// Example:
294 /// ```
295 /// AdaptiveAvgPool3d model(AdaptiveAvgPool3dOptions(3));
296 /// ```
297 using AdaptiveAvgPool3dOptions =
298     AdaptiveAvgPoolOptions<ExpandingArrayWithOptionalElem<3>>;
299 
300 namespace functional {
301 /// Options for `torch::nn::functional::adaptive_avg_pool1d`.
302 ///
303 /// See the documentation for `torch::nn::AdaptiveAvgPool1dOptions` class to
304 /// learn what arguments are supported.
305 ///
306 /// Example:
307 /// ```
308 /// namespace F = torch::nn::functional;
309 /// F::adaptive_avg_pool1d(x, F::AdaptiveAvgPool1dFuncOptions(3));
310 /// ```
311 using AdaptiveAvgPool1dFuncOptions = AdaptiveAvgPool1dOptions;
312 } // namespace functional
313 
314 namespace functional {
315 /// Options for `torch::nn::functional::adaptive_avg_pool2d`.
316 ///
317 /// See the documentation for `torch::nn::AdaptiveAvgPool2dOptions` class to
318 /// learn what arguments are supported.
319 ///
320 /// Example:
321 /// ```
322 /// namespace F = torch::nn::functional;
323 /// F::adaptive_avg_pool2d(x, F::AdaptiveAvgPool2dFuncOptions(3));
324 /// ```
325 using AdaptiveAvgPool2dFuncOptions = AdaptiveAvgPool2dOptions;
326 } // namespace functional
327 
328 namespace functional {
329 /// Options for `torch::nn::functional::adaptive_avg_pool3d`.
330 ///
331 /// See the documentation for `torch::nn::AdaptiveAvgPool3dOptions` class to
332 /// learn what arguments are supported.
333 ///
334 /// Example:
335 /// ```
336 /// namespace F = torch::nn::functional;
337 /// F::adaptive_avg_pool3d(x, F::AdaptiveAvgPool3dFuncOptions(3));
338 /// ```
339 using AdaptiveAvgPool3dFuncOptions = AdaptiveAvgPool3dOptions;
340 } // namespace functional
341 
342 // ============================================================================
343 
344 /// Options for a `D`-dimensional maxunpool module.
345 template <size_t D>
346 struct MaxUnpoolOptions {
MaxUnpoolOptionsMaxUnpoolOptions347   MaxUnpoolOptions(ExpandingArray<D> kernel_size)
348       : kernel_size_(kernel_size), stride_(kernel_size) {}
349 
350   /// the size of the window to take a max over
351   TORCH_ARG(ExpandingArray<D>, kernel_size);
352 
353   /// the stride of the window. Default value is `kernel_size
354   TORCH_ARG(ExpandingArray<D>, stride);
355 
356   /// implicit zero padding to be added on both sides
357   TORCH_ARG(ExpandingArray<D>, padding) = 0;
358 };
359 
360 /// `MaxUnpoolOptions` specialized for the `MaxUnpool1d` module.
361 ///
362 /// Example:
363 /// ```
364 /// MaxUnpool1d model(MaxUnpool1dOptions(3).stride(2).padding(1));
365 /// ```
366 using MaxUnpool1dOptions = MaxUnpoolOptions<1>;
367 
368 /// `MaxUnpoolOptions` specialized for the `MaxUnpool2d` module.
369 ///
370 /// Example:
371 /// ```
372 /// MaxUnpool2d model(MaxUnpool2dOptions(3).stride(2).padding(1));
373 /// ```
374 using MaxUnpool2dOptions = MaxUnpoolOptions<2>;
375 
376 /// `MaxUnpoolOptions` specialized for the `MaxUnpool3d` module.
377 ///
378 /// Example:
379 /// ```
380 /// MaxUnpool3d model(MaxUnpool3dOptions(3).stride(2).padding(1));
381 /// ```
382 using MaxUnpool3dOptions = MaxUnpoolOptions<3>;
383 
384 // ============================================================================
385 
386 namespace functional {
387 
388 /// Options for a `D`-dimensional maxunpool functional.
389 template <size_t D>
390 struct MaxUnpoolFuncOptions {
MaxUnpoolFuncOptionsMaxUnpoolFuncOptions391   MaxUnpoolFuncOptions(ExpandingArray<D> kernel_size)
392       : kernel_size_(kernel_size), stride_(kernel_size) {}
393 
394   /// the size of the window to take a max over
395   TORCH_ARG(ExpandingArray<D>, kernel_size);
396 
397   /// the stride of the window. Default value is `kernel_size
398   TORCH_ARG(ExpandingArray<D>, stride);
399 
400   /// implicit zero padding to be added on both sides
401   TORCH_ARG(ExpandingArray<D>, padding) = 0;
402 
403   /// the targeted output size
404   TORCH_ARG(std::optional<std::vector<int64_t>>, output_size) = std::nullopt;
405 };
406 
407 /// `MaxUnpoolFuncOptions` specialized for
408 /// `torch::nn::functional::max_unpool1d`.
409 ///
410 /// Example:
411 /// ```
412 /// namespace F = torch::nn::functional;
413 /// F::max_unpool1d(x, indices,
414 /// F::MaxUnpool1dFuncOptions(3).stride(2).padding(1));
415 /// ```
416 using MaxUnpool1dFuncOptions = MaxUnpoolFuncOptions<1>;
417 
418 /// `MaxUnpoolFuncOptions` specialized for
419 /// `torch::nn::functional::max_unpool2d`.
420 ///
421 /// Example:
422 /// ```
423 /// namespace F = torch::nn::functional;
424 /// F::max_unpool2d(x, indices,
425 /// F::MaxUnpool2dFuncOptions(3).stride(2).padding(1));
426 /// ```
427 using MaxUnpool2dFuncOptions = MaxUnpoolFuncOptions<2>;
428 
429 /// `MaxUnpoolFuncOptions` specialized for
430 /// `torch::nn::functional::max_unpool3d`.
431 ///
432 /// Example:
433 /// ```
434 /// namespace F = torch::nn::functional;
435 /// F::max_unpool3d(x, indices, F::MaxUnpool3dFuncOptions(3));
436 /// ```
437 using MaxUnpool3dFuncOptions = MaxUnpoolFuncOptions<3>;
438 
439 } // namespace functional
440 
441 // ============================================================================
442 
443 /// Options for a `D`-dimensional fractional maxpool module.
444 template <size_t D>
445 struct FractionalMaxPoolOptions {
FractionalMaxPoolOptionsFractionalMaxPoolOptions446   FractionalMaxPoolOptions(ExpandingArray<D> kernel_size)
447       : kernel_size_(kernel_size) {}
448 
449   /// the size of the window to take a max over
450   TORCH_ARG(ExpandingArray<D>, kernel_size);
451 
452   /// the target output size of the image
453   TORCH_ARG(std::optional<ExpandingArray<D>>, output_size) = std::nullopt;
454 
455   /// If one wants to have an output size as a ratio of the input size, this
456   /// option can be given. This has to be a number or tuple in the range (0, 1)
457   using ExpandingArrayDouble = torch::ExpandingArray<D, double>;
458   TORCH_ARG(std::optional<ExpandingArrayDouble>, output_ratio) = std::nullopt;
459 
460   TORCH_ARG(torch::Tensor, _random_samples) = Tensor();
461 };
462 
463 /// `FractionalMaxPoolOptions` specialized for the `FractionalMaxPool2d` module.
464 ///
465 /// Example:
466 /// ```
467 /// FractionalMaxPool2d model(FractionalMaxPool2dOptions(5).output_size(1));
468 /// ```
469 using FractionalMaxPool2dOptions = FractionalMaxPoolOptions<2>;
470 
471 /// `FractionalMaxPoolOptions` specialized for the `FractionalMaxPool3d` module.
472 ///
473 /// Example:
474 /// ```
475 /// FractionalMaxPool3d model(FractionalMaxPool3dOptions(5).output_size(1));
476 /// ```
477 using FractionalMaxPool3dOptions = FractionalMaxPoolOptions<3>;
478 
479 namespace functional {
480 /// Options for `torch::nn::functional::fractional_max_pool2d` and
481 /// `torch::nn::functional::fractional_max_pool2d_with_indices`
482 ///
483 /// Example:
484 /// ```
485 /// namespace F = torch::nn::functional;
486 /// F::fractional_max_pool2d(x,
487 /// F::FractionalMaxPool2dFuncOptions(3).output_size(2));
488 /// ```
489 using FractionalMaxPool2dFuncOptions = FractionalMaxPool2dOptions;
490 } // namespace functional
491 
492 namespace functional {
493 /// Options for `torch::nn::functional::fractional_max_pool3d` and
494 /// `torch::nn::functional::fractional_max_pool3d_with_indices`
495 ///
496 /// Example:
497 /// ```
498 /// namespace F = torch::nn::functional;
499 /// F::fractional_max_pool3d(x,
500 /// F::FractionalMaxPool3dFuncOptions(3).output_size(2));
501 /// ```
502 using FractionalMaxPool3dFuncOptions = FractionalMaxPool3dOptions;
503 } // namespace functional
504 
505 // ============================================================================
506 
507 /// Options for a `D`-dimensional lppool module.
508 template <size_t D>
509 struct LPPoolOptions {
LPPoolOptionsLPPoolOptions510   LPPoolOptions(double norm_type, ExpandingArray<D> kernel_size)
511       : norm_type_(norm_type),
512         kernel_size_(kernel_size),
513         stride_(kernel_size) {}
514 
515   TORCH_ARG(double, norm_type);
516 
517   // the size of the window to take an average over
518   TORCH_ARG(ExpandingArray<D>, kernel_size);
519 
520   // the stride of the window. Default value is `kernel_size`
521   TORCH_ARG(ExpandingArray<D>, stride);
522 
523   // when True, will use `ceil` instead of `floor` to compute the output shape
524   TORCH_ARG(bool, ceil_mode) = false;
525 };
526 
527 /// `LPPoolOptions` specialized for the `LPPool1d` module.
528 ///
529 /// Example:
530 /// ```
531 /// LPPool1d model(LPPool1dOptions(1, 2).stride(5).ceil_mode(true));
532 /// ```
533 using LPPool1dOptions = LPPoolOptions<1>;
534 
535 /// `LPPoolOptions` specialized for the `LPPool2d` module.
536 ///
537 /// Example:
538 /// ```
539 /// LPPool2d model(LPPool2dOptions(1, std::vector<int64_t>({3, 4})).stride({5,
540 /// 6}).ceil_mode(true));
541 /// ```
542 using LPPool2dOptions = LPPoolOptions<2>;
543 
544 /// `LPPoolOptions` specialized for the `LPPool3d` module.
545 ///
546 /// Example:
547 /// ```
548 /// LPPool3d model(LPPool3dOptions(1, std::vector<int64_t>({3, 4, 5})).stride(
549 /// {5, 6, 7}).ceil_mode(true));
550 /// ```
551 using LPPool3dOptions = LPPoolOptions<3>;
552 
553 namespace functional {
554 /// Options for `torch::nn::functional::lp_pool1d`.
555 ///
556 /// See the documentation for `torch::nn::LPPool1dOptions` class to learn what
557 /// arguments are supported.
558 ///
559 /// Example:
560 /// ```
561 /// namespace F = torch::nn::functional;
562 /// F::lp_pool1d(x, F::LPPool1dFuncOptions(2, 3).stride(2));
563 /// ```
564 using LPPool1dFuncOptions = LPPool1dOptions;
565 } // namespace functional
566 
567 namespace functional {
568 /// Options for `torch::nn::functional::lp_pool2d`.
569 ///
570 /// See the documentation for `torch::nn::LPPool2dOptions` class to learn what
571 /// arguments are supported.
572 ///
573 /// Example:
574 /// ```
575 /// namespace F = torch::nn::functional;
576 /// F::lp_pool2d(x, F::LPPool2dFuncOptions(2, {2, 3}).stride(2));
577 /// ```
578 using LPPool2dFuncOptions = LPPool2dOptions;
579 } // namespace functional
580 
581 namespace functional {
582 /// Options for `torch::nn::functional::lp_pool3d`.
583 ///
584 /// See the documentation for `torch::nn::LPPool3dOptions` class to learn what
585 /// arguments are supported.
586 ///
587 /// Example:
588 /// ```
589 /// namespace F = torch::nn::functional;
590 /// F::lp_pool3d(x, F::LPPool3dFuncOptions(2, {2, 3, 4}).stride(2));
591 /// ```
592 using LPPool3dFuncOptions = LPPool3dOptions;
593 } // namespace functional
594 
595 } // namespace nn
596 } // namespace torch
597