ComfyUI Node

NNT Analyze Model

NntAnalyzeModel — a parameter counter and memory check for your NNT network

By inventorado·Created 2 years ago·Updated 2 years ago· 69
NNT Analyze Model
  • MODEL
  • STRING
input_shape[3, 32, 32]
batch_size1

Three inputs, one text output, zero glamour - and genuinely useful if you're building networks in the Neural Network Toolkit. NntAnalyzeModel takes the model you compiled, pokes through it, and tells you how many parameters it has, how much memory it'll eat, and what's inside. It's the "how big is this thing actually" node.

What it needs

  • MODEL - the model output from NntCompileModel.
  • input_shape - a string, default [3, 32, 32], describing one sample without the batch dimension. That's channels-first: 3 channels, 32×32 for CIFAR-sized images. The node builds a dummy input of torch.randn(batch_size, *shape) to trace shapes through the network, so this has to match what your model actually expects. Feed it the wrong shape and the whole thing is just a guess.
  • batch_size - default 1; mostly affects the dummy input and the flavor of the report, not the math of parameter counts.

Output: a single STRING - the report. Wire it into a text-display node (the pack's example workflows use JjkShowText for this, via ComfyUI-Jjk-Nodes) or just read it from the node's output.

What you get back

The report is a plain-text breakdown:

  • Total and trainable parameters, counted by walking named_modules() and summing numel() over Linear, Conv2d, and BatchNorm2d layers, with every other module listed as an activation.
  • Memory estimate - parameter bytes plus buffer bytes, converted to MB. Parameters-only math, so treat it as a lower bound; activations and gradients are where real training memory lives.
  • Layer-by-layer architecture - each layer with its in/out features, kernel, stride, padding.

For a teaching tool this is a lovely feature: it's the "show your work" node. You can watch your parameter count blow up as you add dense layers, and start to feel why a 1024-node dense layer after a flatten is a terrible idea on a small dataset.

Honest caveats

The counting is a bit loose. It special-cases Linear, Conv2d, and BatchNorm2d - a Conv1d or GroupNorm will get reported as an activation rather than counted as parameters. And since it moves the model to CPU and reports the architecture without a forward pass, "memory" is purely a weight-count, not what training will actually use. Fine for eyeballing, not for capacity planning.

Also - the obvious trap - this analyzes the little NNT model you compiled, not any diffusion checkpoint. If you're searching for "analyze model" hoping for a report on your SDXL or Flux model, this isn't the node; that's a totally different thing.

Install

Comes with the pack, no per-node install:

cd ComfyUI/custom_nodes
git clone https://github.com/inventorado/ComfyUI_NNT.git
cd ComfyUI_NNT
pip install -r requirements.txt

or find "ComfyUI Neural Network Toolkit NNT" in ComfyUI Manager. Since it pulls in torchview and graphviz among its dependencies, be ready for a longer first install. Restart ComfyUI after, then drop it at the end of any compile workflow and see exactly what you built.

CategoryNNT Neural Network Toolkit/Models

Inputs (3)

NameTypeDefaultDescription
MODELMODEL
input_shapeSTRING[3, 32, 32]
batch_sizeINT11–128

Outputs (1)

NameTypeDescription
STRINGSTRING