conv2d_module.hpp

Go to the documentation of this file.

 
#pragma once
 
#include <optional>
#include <initializer_list>
#include <vector>
 
#include "pulsatrix/module.hpp"
 
namespace pulsatrix {
 
class Conv2DModule : public Module {
public:
    Conv2DModule(int64_t in_channels, int64_t out_channels, int64_t kernel_h, int64_t kernel_w,
                 DeviceBackend* backend, int64_t stride = 1, int64_t padding = 0);
 
    [[nodiscard]] int64_t stride() const { return stride_; }
    [[nodiscard]] int64_t padding() const { return padding_; }
 
    [[nodiscard]] Tensor backward(const Tensor& grad_output) override;
 
    [[nodiscard]] OpType op_type() const override { return OpType::Conv; }
 
    void set_kernel(std::initializer_list<float> values);
 
    void set_bias(std::initializer_list<float> values);
 
    void set_kernel(const std::vector<float>& values);
 
    void set_bias(const std::vector<float>& values);
 
    [[nodiscard]] const Tensor& kernel() const { return kernel_; }
    [[nodiscard]] const Tensor& bias() const { return bias_; }
    [[nodiscard]] const Tensor& kernel_grad() const { return kernel_grad_; }
    [[nodiscard]] const Tensor& bias_grad() const { return bias_grad_; }
 
    [[nodiscard]] Tensor propagate_relevance(const Tensor& relevance_out, const LRPRuleConfig& config) override;
 
    [[nodiscard]] bool supports_lrp_rule(LRPRule) const override { return true; }
 
    [[nodiscard]] std::vector<NamedParamRef> named_parameters() override {
        return {{"weight", {&kernel_, &kernel_grad_}}, {"bias", {&bias_, &bias_grad_}}};
    }
 
 
    [[nodiscard]] std::optional<DeviceType> compute_device() const override { return backend_->device(); }
 
protected:
    [[nodiscard]] Tensor forward_impl(const Tensor& input) override;
 
private:
    int64_t in_channels_;
    int64_t out_channels_;
    int64_t kernel_h_;
    int64_t kernel_w_;
    int64_t stride_;
    int64_t padding_;
    DeviceBackend* backend_;
    Tensor kernel_;       // shape (out_channels, in_channels, kernel_h, kernel_w)
    Tensor bias_;         // shape (out_channels,)
    Tensor kernel_grad_;
    Tensor bias_grad_;
    Tensor last_input_;       // (N, in_channels, H, W)
    Tensor last_im2col_;      // (N, patch_size, out_h*out_w), cached for backward
    Tensor last_pre_bias_output_;  // (N, out_channels, out_h, out_w), cached for LRP
    int64_t last_out_h_ = 0;
    int64_t last_out_w_ = 0;
    bool has_forwarded_ = false;
 
    [[nodiscard]] ConvGeometry geometry() const {
        const auto s = static_cast<size_t>(stride_), p = static_cast<size_t>(padding_);
        return {static_cast<size_t>(kernel_h_), static_cast<size_t>(kernel_w_), s, s, p, p};
    }
};
 
}  // namespace pulsatrix