tinytensor/include/tt/nn/pool2d.h at master · tuero/tinytensor · GitHub

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
// pool2d.h
// Pooling layer modules

#ifndef TINYTENSOR_NN_POOL2D_H_
#define TINYTENSOR_NN_POOL2D_H_

#include <tt/export.h>
#include <tt/nn/module.h>
#include <tt/tensor.h>

#include <ostream>
#include <string>

namespace tinytensor::nn {

// A min pooling layer
class TINYTENSOR_EXPORT MinPool2d : public Module {
public:
    /**
     * Construct a MinPool2d layer
     * @param kernel_size The kernel size
     * @param stride The stride
     * @param padding The amount of padding to apply to each side of the input
     */
    MinPool2d(int kernel_size, int stride, int padding);

    [[nodiscard]] auto forward(const Tensor &input) const -> Tensor;

    void pretty_print(std::ostream &os) const override;

    [[nodiscard]] auto name() const -> std::string override {
        return "MinPool2D";
    }

private:
    int kernel_size_;
    int stride_;
    int padding_;
};

// A max pooling layer
class TINYTENSOR_EXPORT MaxPool2d : public Module {
public:
    /**
     * Construct a MaxPool2d layer
     * @param kernel_size The kernel size
     * @param stride The stride
     * @param padding The amount of padding to apply to each side of the input
     */
    MaxPool2d(int kernel_size, int stride, int padding);

    [[nodiscard]] auto forward(const Tensor &input) const -> Tensor;

    void pretty_print(std::ostream &os) const override;

    [[nodiscard]] auto name() const -> std::string override {
        return "MaxPool2D";
    }

private:
    int kernel_size_;
    int stride_;
    int padding_;
};

// A average pooling layer
class TINYTENSOR_EXPORT AvgPool2d : public Module {
public:
    /**
     * Construct a AvgPool2d layer
     * @param kernel_size The kernel size
     * @param stride The stride
     * @param padding The amount of padding to apply to each side of the input
     */
    AvgPool2d(int kernel_size, int stride, int padding);

    [[nodiscard]] auto forward(const Tensor &input) const -> Tensor;

    void pretty_print(std::ostream &os) const override;

    [[nodiscard]] auto name() const -> std::string override {
        return "AvgPool2D";
    }

private:
    int kernel_size_;
    int stride_;
    int padding_;
};

}    // namespace tinytensor::nn

#endif    // TINYTENSOR_NN_POOL2D_H_