NNT Edit Model Layers

Once you've trained a model in the Neural Network Toolkit, the fun - and the fiddly parts - begin. NNT Edit Model Layers is the post-training surgery node: freeze layers for transfer learning, reinitialize weights, prune the dead ones, or quantize. It's the "tune a finished model" tool in a pack that's otherwise all about building from scratch, and for an educational toolkit it's a great way to see what these operations actually do to your weights.

What it actually does

It takes a compiled MODEL in, walks its named modules, and applies one of five operations to a selected subset of layers:

• Freeze / Unfreeze - flips requires_grad on the chosen parameters. Freezing the feature-extractor layers and only training the head is the classic transfer-learning move; this node makes that a dropdown selection instead of a for loop.
• Set weights and biases - reinitializes parameters with any of nine schemes: random, zeros, ones, xavier_uniform/normal, kaiming_uniform/normal, orthogonal, or custom_value. The pack's author specifically recommends kaiming_normal over xavier_normal for ReLU-based networks - worth remembering.
• Prune - magnitude pruning: zeroes out the smallest pruning_amount (as a fraction, 0–0.9) of weights below a threshold. Warning: it's an unstructured zero-out, not a speedup - the model gets sparser but not smaller.
• Quantize - rounds weights to quantization_bits with a naive min/max scale. This is teaching-grade quantization, not calibrated PTQ; it'll show you the concept and the accuracy cost, but treat the result as a lesson, not a production model.

Inputs that matter

• operation - Freeze, Unfreeze, Set weights and biases, Prune, or Quantize. Everything else on the node exists to scope this.
• layer_selection - All layers, First N layers, Last N layers, or Selected types (with layer_types filtering to Linear / Conv2d / BatchNorm2d / All). The head is usually the "last N" - freeze everything else.
• parameter_type - weights, biases, or both.
• num_layers - how many layers "First N"/"Last N" means.

Then the operation-specific knobs: initialization + custom_value for the reinit operation, pruning_amount, quantization_bits.

Outputs

Three: the edited_model (wire it into training or inference), an info_message string summarizing what changed, and layer_stats - a DICT with total/trainable parameter counts, a per-layer-type histogram, and min/max/mean/std for every parameter tensor. The stats dict is genuinely useful for checking before you fine-tune that the right layers are frozen.

Installing NNT

Part of inventorado/ComfyUI_NNT. ComfyUI Manager (search "ComfyUI Neural Network Toolkit") or:

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

Restart ComfyUI after. Heavy requirements - torch, scikit-learn, pandas, transformers, shap - so budget the first install. The pack's example workflows want ComfyUI-Jjk-Nodes for text output; let Manager fetch it. And remember the README's honest framing: this is a learning/prototyping toolkit, and the quantization in particular is about seeing how the idea works, not shipping a 4-bit mobile model.