grad_cam.hpp

Go to the documentation of this file.

 
#pragma once
 
#include <stdexcept>
#include <string>
#include <vector>
 
#include "pulsatrix/assert.hpp"
#include "pulsatrix/attribution.hpp"
#include "pulsatrix/device_backend.hpp"
#include "pulsatrix/explainer_context.hpp"
#include "pulsatrix/tensor.hpp"
 
namespace pulsatrix {
 
class GradCAM {
public:
    [[nodiscard]] Attribution explain(ExplainerContext& ctx, const Tensor& input, int64_t target_index,
                                       DeviceBackend* backend) const {
        Tensor output = ctx.forward_pass(input);
        // External boundary (Mission 2, finding 15 systemic sweep; rank check added by
        // campaign_exai_dl_library_batch_dimension_support) -- escalated from
        // PULSATRIX_ASSERT-only.
        if (output.rank() != 2) {
            throw std::invalid_argument("GradCAM::explain: network output must be rank-2 (N, num_classes)");
        }
        if (target_index < 0 || target_index >= output.shape().dim(1)) {
            throw std::invalid_argument("GradCAM::explain: target_index out of range");
        }
        int64_t N = output.shape().dim(0);
        int64_t num_classes = output.shape().dim(1);
 
        std::vector<float> seed_values(static_cast<size_t>(output.numel()), 0.0f);
        for (int64_t n = 0; n < N; ++n) {
            seed_values[static_cast<size_t>(n * num_classes + target_index)] = 1.0f;
        }
        Tensor seed(output.shape(), explainer_detail::backend_beside(output, backend), seed_values,
                    output.device());
        (void)ctx.backward_pass(seed);
 
        std::vector<NodeId> conv_nodes = ctx.graph().nodes_by_op_type(OpType::Conv);
        // External boundary -- whether this ExplainerContext was built with a Conv layer
        // is a caller-configuration fact, not an internal invariant this library controls.
        if (conv_nodes.empty()) {
            throw std::invalid_argument("GradCAM::explain: graph has no Conv layer");
        }
        NodeId target_node = conv_nodes.back();
 
        const Tensor& activation_tensor = ctx.activation(target_node);
        const Tensor& grad_tensor = ctx.gradient(target_node);
 
        int64_t channels = activation_tensor.shape().dim(1);
        int64_t height = activation_tensor.shape().dim(2);
        int64_t width = activation_tensor.shape().dim(3);
 
        // Host boundary: one read of each buffer, then the original arithmetic on the copies.
        const std::vector<float> activation = activation_tensor.to_host_vector();
        const std::vector<float> grad = grad_tensor.to_host_vector();
        auto at4 = [&](const std::vector<float>& v, int64_t n, int64_t c, int64_t h, int64_t w) {
            return v[static_cast<size_t>(((n * channels + c) * height + h) * width + w)];
        };
 
        std::vector<float> cam(static_cast<size_t>(N * height * width), 0.0f);
        for (int64_t n = 0; n < N; ++n) {
            std::vector<float> alpha(static_cast<size_t>(channels), 0.0f);
            for (int64_t c = 0; c < channels; ++c) {
                float sum = 0.0f;
                for (int64_t h = 0; h < height; ++h) {
                    for (int64_t w = 0; w < width; ++w) {
                        sum += at4(grad, n, c, h, w);
                    }
                }
                alpha[static_cast<size_t>(c)] = sum / static_cast<float>(height * width);
            }
 
            for (int64_t h = 0; h < height; ++h) {
                for (int64_t w = 0; w < width; ++w) {
                    float value = 0.0f;
                    for (int64_t c = 0; c < channels; ++c) {
                        value += alpha[static_cast<size_t>(c)] * at4(activation, n, c, h, w);
                    }
                    cam[static_cast<size_t>((n * height + h) * width + w)] = value > 0.0f ? value : 0.0f;
                }
            }
        }
 
        Tensor cam_tensor(Shape({N, height, width}), explainer_detail::backend_beside(activation_tensor, backend),
                          cam, activation_tensor.device());
        return Attribution{"grad_cam", std::move(cam_tensor),
                            {{"target_index", std::to_string(target_index)},
                             {"layer_node_id", std::to_string(target_node)}}};
    }
};
 
}  // namespace pulsatrix