NNT SHAP Summary Node
Find out which features your toy model actually listens to
- model
- X_train_sample
- X_test_sample
- text_report
- shap_plot_tensor
Your model got to 90% accuracy. Cool. But why? NNT SHAP Summary Node is the explainability tool of the Neural Network Toolkit - it runs SHAP (SHapley Additive exPlanations) on your trained model and tells you which input features are actually driving its predictions. If you're using this pack to learn ML, this is the node that makes "the model learned the signal, not the noise" feel tangible instead of abstract.
The tradeoff is worth naming up front: SHAP is expensive. This node uses shap.KernelExplainer, a model-agnostic method that approximates Shapley values by sampling. On a small toy dataset it's perfectly fine. On anything real-sized it will take forever. Keep your samples small and your expectations calibrated.
How it works
You feed it a trained MODEL plus two tensors: X_train_sample (a slice of your training data) and X_test_sample (the rows you want explained). It takes the first background_sample_size rows of the training sample (default 100) as the "background" that SHAP uses as a baseline, then explains each test row against it. Choose a plot_type - dot, bar, or violin - and it produces:
- text_report - a string listing the top 10 features by mean absolute SHAP value.
- shap_plot_tensor - the summary plot rendered as a
TENSOR(a raw CHW torch tensor, not a ComfyUIIMAGE).
That last bit is the subtle gotcha. To actually see the plot, you route shap_plot_tensor through the pack's NntTensorElementToImage node, which converts a tensor element into a viewable image. Don't try to wire it straight into a Save Image node - the types don't match, and ComfyUI will just shrug at you.
There's another quirk baked in: the plotting code indexes shap_values[1], which assumes a binary classifier where class 1 is the positive class. If your model has more than two classes or your SHAP values come back shaped differently, the plot can throw or come out wrong even though the text report is fine. It's an educational toolkit, and this is one of those places where the "educational" cuts both ways.
Install and usage
Install with the pack:
cd ComfyUI/custom_nodes
git clone https://github.com/inventorado/ComfyUI_NNT.git
cd ComfyUI_NNT
pip install -r requirements.txt
Restart ComfyUI (or install via Manager, searching "ComfyUI Neural Network Toolkit NNT"). The requirements pin shap==0.41.0 and pull in numba, cloudpickle, and slicer for it - if you already have a newer numba in your environment, watch for conflicts, since shap's old version can be picky.
Practical tip: keep background_sample_size modest. If it exceeds the number of rows in X_train_sample, the node throws a clear error telling you so. And since KernelExplainer is the slow path, start with a few dozen background rows and a small test sample, sanity-check the report, then scale up. The text report alone is often enough to learn from - the plot is the reward for patience.
Inputs (5)
| Name | Type | Default | Description |
|---|---|---|---|
| model | MODEL | — | |
| X_train_sample | TENSOR | — | |
| X_test_sample | TENSOR | — | |
| plot_type | COMBO | dot | 3 options: dot, bar, violin |
| background_sample_sizeopt | INT | 1001–10000 | — |
Outputs (2)
| Name | Type | Description |
|---|---|---|
| text_report | STRING | — |
| shap_plot_tensor | TENSOR | — |