batch_norm_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 BatchNormFold;
 
class BatchNormModule : public Module {
public:
    BatchNormModule(int64_t num_channels, DeviceBackend* backend, DeviceType device,
                     float eps = 1e-6f, float momentum = 0.1f);
 
    BatchNormModule(int64_t num_channels, DeviceBackend* backend);
 
    [[nodiscard]] Tensor backward(const Tensor& grad_output) override;
 
    [[nodiscard]] OpType op_type() const override { return OpType::Normalization; }
 
    void set_gamma(std::initializer_list<float> values);
    void set_beta(std::initializer_list<float> values);
    void set_gamma(const std::vector<float>& values);
    void set_beta(const std::vector<float>& values);
 
    [[nodiscard]] const Tensor& running_mean() const { return running_mean_; }
    [[nodiscard]] const Tensor& running_var() const { return running_var_; }
    void set_running_mean(const std::vector<float>& values);
    void set_running_var(const std::vector<float>& values);
 
    [[nodiscard]] const Tensor& gamma() const { return gamma_; }
    [[nodiscard]] const Tensor& beta() const { return beta_; }
    [[nodiscard]] const Tensor& gamma_grad() const { return gamma_grad_; }
    [[nodiscard]] const Tensor& beta_grad() const { return beta_grad_; }
 
    [[nodiscard]] Tensor propagate_relevance(const Tensor& relevance_out, const LRPRuleConfig& config) override;
 
    [[nodiscard]] std::vector<NamedParamRef> named_parameters() override {
        return {{"weight", {&gamma_, &gamma_grad_}}, {"bias", {&beta_, &beta_grad_}}};
    }
 
 
    [[nodiscard]] std::optional<DeviceType> compute_device() const override { return gamma_.device(); }
 
protected:
    [[nodiscard]] Tensor forward_impl(const Tensor& input) override;
 
private:
    friend class BatchNormFold;
 
    // Which statistics the cached forward used, so backward() matches it even if the mode
    // changes in between.
    enum class Mode { Batch, Running, Folded };
 
    int64_t num_channels_;
    float eps_;
    float momentum_;
    DeviceBackend* backend_;
    Tensor gamma_;       // shape (num_channels,)
    Tensor beta_;        // shape (num_channels,)
    Tensor gamma_grad_;
    Tensor beta_grad_;
    Tensor last_input_;      // (N, num_channels, H, W)
    Tensor last_xhat_;       // (N, num_channels, H, W)
    Tensor last_std_;  // (num_channels,) one std per channel, on the module's device
    Tensor running_mean_;  // (num_channels,)
    Tensor running_var_;   // (num_channels,)
    bool has_forwarded_ = false;
    Mode last_mode_ = Mode::Batch;
    bool folded_ = false;  // set by BatchNormFold: forward/backward become the identity
};
 
}  // namespace pulsatrix