lrp.hpp

Go to the documentation of this file.

 
#pragma once
 
#include <cstdint>
#include <functional>
#include <memory>
#include <stdexcept>
#include <string>
#include <utility>
#include <vector>
 
#include "pulsatrix/attribution.hpp"
#include "pulsatrix/conv2d_module.hpp"
#include "pulsatrix/device_backend.hpp"
#include "pulsatrix/explainer_context.hpp"
#include "pulsatrix/linear_module.hpp"
#include "pulsatrix/lrp_rule_config.hpp"
#include "pulsatrix/tensor.hpp"
 
namespace pulsatrix {
 
enum class LRPSeed {
    OutputValue,
    OneHot
};
 
struct LRPTarget {
    std::vector<int64_t> targets;
    std::vector<int64_t> contrasts = {};
    LRPSeed seed = LRPSeed::OutputValue;
};
 
using LRPComposite = std::function<LRPRuleConfig(size_t layer_index, const Module& module)>;
 
namespace lrp_composite {
 
namespace detail {
inline bool is_conv(const Module& module) { return dynamic_cast<const Conv2DModule*>(&module) != nullptr; }
inline LRPRuleConfig zennit_epsilon(float epsilon) {
    LRPRuleConfig config{epsilon};
    config.epsilon_bias_in_denominator = true;
    return config;
}
inline LRPRuleConfig alpha_beta(float alpha, float beta, float epsilon) {
    LRPRuleConfig config{epsilon};
    config.rule = LRPRule::AlphaBeta;
    config.alpha = alpha;
    config.beta = beta;
    return config;
}
}  // namespace detail
 
inline LRPComposite epsilon_plus(float epsilon = 1e-6f) {
    return [epsilon](size_t, const Module& module) {
        return detail::is_conv(module) ? detail::alpha_beta(1.0f, 0.0f, epsilon) : detail::zennit_epsilon(epsilon);
    };
}
 
inline LRPComposite epsilon_alpha2_beta1(float epsilon = 1e-6f) {
    return [epsilon](size_t, const Module& module) {
        return detail::is_conv(module) ? detail::alpha_beta(2.0f, 1.0f, epsilon) : detail::zennit_epsilon(epsilon);
    };
}
 
inline LRPComposite epsilon_gamma_box(float low, float high, float gamma = 0.25f, float epsilon = 1e-6f) {
    auto first_conv_seen = std::make_shared<bool>(false);
    return [=](size_t layer_index, const Module& module) {
        if (layer_index == 0) {
            *first_conv_seen = false;
        }
        if (!detail::is_conv(module)) {
            return detail::zennit_epsilon(epsilon);
        }
        LRPRuleConfig config{epsilon};
        if (!*first_conv_seen) {
            *first_conv_seen = true;
            config.rule = LRPRule::ZBox;
            config.low = low;
            config.high = high;
        } else {
            config.rule = LRPRule::Gamma;
            config.gamma = gamma;
        }
        return config;
    };
}
 
}  // namespace lrp_composite
 
class LRP {
public:
    explicit LRP(LRPRuleConfig config = LRPRuleConfig{}) : config_(config) {}
 
    explicit LRP(LRPComposite composite, std::string name = "custom")
        : config_(), composite_(std::move(composite)), composite_name_(std::move(name)) {
        if (!composite_) {
            throw std::invalid_argument("LRP: composite must not be empty");
        }
    }
 
    [[nodiscard]] static LRP epsilon_plus(float epsilon = 1e-6f) {
        return LRP(lrp_composite::epsilon_plus(epsilon), "epsilon_plus");
    }
    [[nodiscard]] static LRP epsilon_gamma_box(float low, float high, float gamma = 0.25f, float epsilon = 1e-6f) {
        return LRP(lrp_composite::epsilon_gamma_box(low, high, gamma, epsilon), "epsilon_gamma_box");
    }
    [[nodiscard]] static LRP epsilon_alpha2_beta1(float epsilon = 1e-6f) {
        return LRP(lrp_composite::epsilon_alpha2_beta1(epsilon), "epsilon_alpha2_beta1");
    }
 
    [[nodiscard]] Attribution explain(ExplainerContext& ctx, const Tensor& input, int64_t target_index,
                                      DeviceBackend* backend) const {
        return explain(ctx, input, LRPTarget{{target_index}}, backend);
    }
 
    [[nodiscard]] Attribution explain(ExplainerContext& ctx, const Tensor& input, const LRPTarget& target,
                                      DeviceBackend* backend) const {
        Tensor output = ctx.forward_pass(input);
        if (output.rank() != 2) {
            throw std::invalid_argument("LRP::explain: network output must be rank-2 (N, num_classes)");
        }
        const int64_t N = output.shape().dim(0);
        const int64_t C = output.shape().dim(1);
        const std::vector<int64_t> targets = per_row(target.targets, N, C, "targets");
        const std::vector<int64_t> contrasts =
            target.contrasts.empty() ? std::vector<int64_t>{} : per_row(target.contrasts, N, C, "contrasts");
        for (size_t n = 0; n < contrasts.size(); ++n) {
            if (contrasts[n] == targets[n]) {
                throw std::invalid_argument(
                    "LRP::explain: a row's contrast equals its target (the seed would be zero)");
            }
        }
 
        // Seed on the host, then upload: valid whatever device the output lives on.
        std::vector<float> y;
        if (target.seed == LRPSeed::OutputValue) {
            y.resize(static_cast<size_t>(N * C));
            output.backend()->copy(y.data(), output.data(), y.size() * sizeof(float),
                                   output.device() == DeviceType::Cpu ? CopyDirection::HostToHost
                                                                      : CopyDirection::DeviceToHost);
        }
        auto seed_value = [&](int64_t n, int64_t c) {
            return target.seed == LRPSeed::OutputValue ? y[static_cast<size_t>(n * C + c)] : 1.0f;
        };
        std::vector<float> seed(static_cast<size_t>(N * C), 0.0f);
        float relevance_out_sum = 0.0f;
        for (int64_t n = 0; n < N; ++n) {
            const int64_t t = targets[static_cast<size_t>(n)];
            seed[static_cast<size_t>(n * C + t)] = seed_value(n, t);
            relevance_out_sum += seed_value(n, t);
            if (!contrasts.empty()) {
                const int64_t c = contrasts[static_cast<size_t>(n)];
                seed[static_cast<size_t>(n * C + c)] = -seed_value(n, c);
                relevance_out_sum -= seed_value(n, c);
            }
        }
        Tensor seed_tensor(output.shape(), backend, seed, output.device());
 
        std::vector<LRPRuleConfig> configs;
        configs.reserve(ctx.modules().size());
        for (size_t i = 0; i < ctx.modules().size(); ++i) {
            configs.push_back(composite_ ? composite_(i, *ctx.modules()[i]) : config_);
        }
        std::string rules;
        for (size_t i = 0; i < configs.size(); ++i) {
            rules += (i == 0 ? "" : ",") + lrp_rule_name(configs[i].rule);
        }
        Tensor relevance = ctx.relevance_pass(seed_tensor, configs);
        const float relevance_in_sum =
            relevance.backend()->sum(relevance.data(), static_cast<size_t>(relevance.numel()));
 
        return Attribution{"lrp",
                           std::move(relevance),
                           {{"rule", composite_ ? "composite:" + composite_name_ : lrp_rule_name(config_.rule)},
                            {"rules", rules},
                            {"epsilon", std::to_string(config_.epsilon)},
                            {"seed", target.seed == LRPSeed::OutputValue ? "output_value" : "one_hot"},
                            {"targets", join(target.targets)},
                            {"contrasts", join(target.contrasts)},
                            {"relevance_out_sum", std::to_string(relevance_out_sum)},
                            {"relevance_in_sum", std::to_string(relevance_in_sum)}}};
    }
 
    [[nodiscard]] const LRPRuleConfig& config() const { return config_; }
 
    [[nodiscard]] const LRPComposite& composite() const { return composite_; }
 
private:
    // Expands a 1-entry list to N rows and validates every index against C classes.
    static std::vector<int64_t> per_row(const std::vector<int64_t>& indices, int64_t N, int64_t C,
                                        const char* what) {
        if (indices.size() != 1 && static_cast<int64_t>(indices.size()) != N) {
            throw std::invalid_argument(std::string("LRP::explain: ") + what +
                                        " must have 1 entry or one per row of the output");
        }
        std::vector<int64_t> rows(static_cast<size_t>(N));
        for (int64_t n = 0; n < N; ++n) {
            const int64_t index = indices.size() == 1 ? indices[0] : indices[static_cast<size_t>(n)];
            if (index < 0 || index >= C) {
                throw std::invalid_argument(std::string("LRP::explain: ") + what + " index out of range");
            }
            rows[static_cast<size_t>(n)] = index;
        }
        return rows;
    }
 
    static std::string join(const std::vector<int64_t>& v) {
        std::string out;
        for (size_t i = 0; i < v.size(); ++i) {
            out += (i == 0 ? "" : ",") + std::to_string(v[i]);
        }
        return out;
    }
 
    LRPRuleConfig config_;
    LRPComposite composite_;
    std::string composite_name_;
};
 
}  // namespace pulsatrix