max_pool2d_module.hpp

Go to the documentation of this file.

 
#pragma once
 
#include <optional>
#include <vector>
 
#include "pulsatrix/module.hpp"
 
namespace pulsatrix {
 
class MaxPool2DModule : public Module {
public:
    MaxPool2DModule(int64_t kernel_h, int64_t kernel_w, DeviceBackend* backend);
 
    [[nodiscard]] Tensor backward(const Tensor& grad_output) override;
 
    [[nodiscard]] OpType op_type() const override { return OpType::Pooling; }
 
    [[nodiscard]] Tensor propagate_relevance(const Tensor& relevance_out, const LRPRuleConfig& config) override;
 
    [[nodiscard]] bool supports_lrp_rule(LRPRule) const override { return true; }
 
 
    [[nodiscard]] std::optional<DeviceType> compute_device() const override { return backend_->device(); }
 
protected:
    [[nodiscard]] Tensor forward_impl(const Tensor& input) override;
 
private:
    int64_t kernel_h_;
    int64_t kernel_w_;
    DeviceBackend* backend_;
    Shape last_input_shape_ = Shape({0});
    int64_t last_out_h_ = 0;
    int64_t last_out_w_ = 0;
    // Flat (within-window) offset of the argmax for each output element, indexed the same
    // way as the output buffer (n, c, oh, ow) -- row-major flat index into (H, W) input
    // plane, i.e. ih * W + iw. Sized N*C*out_h*out_w after a real forward() call.
    // Stored as whole-number floats on the input's device (exact below 2^24 -- far above any
    // plane this module pools), consumed by DeviceBackend::max_unpool (GPU-native-kernels
    // Mission 4).
    Tensor argmax_flat_index_ = Tensor(Shape({0}), backend_);
    bool has_forwarded_ = false;
};
 
}  // namespace pulsatrix