module.hpp

Go to the documentation of this file.

 
#pragma once
 
#include <optional>
#include <stdexcept>
#include <string>
#include <utility>
#include <vector>
 
#include "pulsatrix/assert.hpp"
#include "pulsatrix/autograd.hpp"
#include "pulsatrix/computation_graph.hpp"
#include "pulsatrix/lrp_rule_config.hpp"
#include "pulsatrix/node.hpp"
#include "pulsatrix/op_type.hpp"
#include "pulsatrix/tensor.hpp"
 
namespace pulsatrix {
 
struct ParamRef {
    Tensor* value;
    Tensor* grad;
};
 
struct NamedParamRef {
    std::string name;
    ParamRef ref;
};
 
class Module {
public:
    virtual ~Module() = default;
 
    [[nodiscard]] Tensor forward(const Tensor& input) {
        if (input.numel() <= 0) {
            throw std::invalid_argument("Module::forward: input must not be empty");
        }
        if (const std::optional<DeviceType> device = compute_device()) {
            require_device(input, *device, "Module::forward");
        }
        return forward_impl(input);
    }
 
    [[nodiscard]] virtual std::optional<DeviceType> compute_device() const { return std::nullopt; }
 
    [[nodiscard]] std::pair<Tensor, NodeId> forward_traced(const Tensor& input, NodeId input_node,
                                                             ComputationGraph& graph, Autograd& autograd) {
        Tensor output = forward(input);
        NodeId node_id = graph.add_node(op_type(), output.shape(), std::nullopt, {input_node});
        autograd.register_backward(node_id, [this](const Tensor& grad_output) { return this->backward(grad_output); });
        return {std::move(output), node_id};
    }
 
    [[nodiscard]] virtual Tensor propagate_relevance(const Tensor& relevance_out, const LRPRuleConfig& config) = 0;
 
    [[nodiscard]] virtual bool supports_lrp_rule(LRPRule rule) const { return rule == LRPRule::Epsilon; }
 
    [[nodiscard]] virtual Tensor backward(const Tensor& grad_output) = 0;
 
    [[nodiscard]] virtual OpType op_type() const = 0;
 
    [[nodiscard]] virtual std::vector<NamedParamRef> named_parameters() { return {}; }
 
    [[nodiscard]] virtual std::vector<ParamRef> parameters() {
        std::vector<ParamRef> params;
        for (const NamedParamRef& p : named_parameters()) {
            params.push_back(p.ref);
        }
        return params;
    }
 
    void set_requires_grad(bool requires_grad, const std::string& prefix = "") {
        if (prefix.empty()) {
            // parameters(), not named_parameters(): also reaches a legacy module that has no names.
            for (ParamRef p : parameters()) {
                p.value->set_requires_grad(requires_grad);
            }
            return;
        }
        std::vector<Tensor*> selected;
        for (const NamedParamRef& p : named_parameters()) {
            if (p.name == prefix || p.name.rfind(prefix + ".", 0) == 0) {
                selected.push_back(p.ref.value);
            }
        }
        if (selected.empty()) {
            throw std::invalid_argument("Module::set_requires_grad: no parameter named or under '" + prefix + "'");
        }
        for (Tensor* value : selected) {
            value->set_requires_grad(requires_grad);
        }
    }
 
    virtual void set_training(bool training) { training_ = training; }
 
    [[nodiscard]] bool is_training() const { return training_; }
 
protected:
    [[nodiscard]] virtual Tensor forward_impl(const Tensor& input) = 0;
 
private:
    bool training_ = true;
};
 
inline void append_named_parameters(std::vector<NamedParamRef>& out, const std::string& prefix, Module& child) {
    std::vector<NamedParamRef> named = child.named_parameters();
    if (named.empty()) {
        std::vector<ParamRef> plain = child.parameters();
        for (size_t i = 0; i < plain.size(); ++i) {
            named.push_back({std::to_string(i), plain[i]});
        }
    }
    for (NamedParamRef& p : named) {
        out.push_back({prefix + "." + p.name, p.ref});
    }
}
 
}  // namespace pulsatrix