Nodes/ComfyUI Neural Network Toolkit NNT /NNT Define Normalization Layer
ComfyUI Node

NNT Define Normalization Layer

The node your loss curve will thank you for

By inventorado·Created 2 years ago·Updated 2 years ago· 69
NNT Define Normalization Layer
  • LAYER_STACK
  • LAYER_STACK
norm_typeBatchNorm2d
num_features64
eps0.0000
momentum0.1
affineTrue
track_running_statsTrue

If your NNT model trains slowly, refuses to converge, or spits out NaNs after a few epochs, the fix is often one node: NNT Define Normalization Layer. It slots a normalization layer into your layer stack - batch, layer, instance, group, or local-response norm - and it's the single highest-leverage addition you can make to a deep network that's being fussy. Normalization re-centers and re-scales activations mid-network, which lets you raise the learning rate and typically cuts epochs to convergence by a lot.

What it actually does

Like every define node in the Neural Network Toolkit, it appends a layer description dict to a LAYER_STACK list. When that stack reaches NntCompileModel, the norm entry is turned into the matching PyTorch module - BatchNorm2d, LayerNorm, and so on - and inserted into the compiled nn.Sequential. The LAYER_STACK input on the node is optional: leave it empty to start fresh, or wire in the previous node's output to keep building.

Inputs that matter

  • norm_type - the dropdown lists None, BatchNorm, LayerNorm, InstanceNorm, GroupNorm, and LocalResponseNorm. Quirk: the default value on the widget is BatchNorm2d, which isn't in that dropdown list at all. It still compiles correctly (the compile step recognizes the concrete BatchNorm2d), so you can leave it as-is for 2D image work - just don't be confused when the value and the menu disagree.
  • num_features - the number of channels (for batch/instance/group norm) or the feature dimension. Set this to the output channels of the layer feeding the norm. Wrong value = runtime error at compile time, which is NNT's way of telling you to count your conv channels again.
  • eps - a tiny stability constant added inside the normalization; the default 1e-05 is fine for almost everything. Raise it if you hit numerical blowups in extreme low-precision situations.

Then momentum (how fast running statistics update for batch norm), affine (whether the layer learns its own scale/shift - leave True), and track_running_stats (batch norm bookkeeping; True for training, and note that False turns batch norm into something closer to instance norm). For a first model, change nothing but num_features and go.

Gotchas

Batch norm is the one that bites beginners. It behaves differently in train vs. eval mode - it uses batch statistics while training and the running averages afterward - and it wants a batch size bigger than 1, so a "batch" of a single sample will behave oddly. If you're doing a tiny NNT toy model with batch size 1, LayerNorm or InstanceNorm is the safer pick. The other classic trap is putting batch norm right before the output layer on a classification head; keep it inside the body of the network and you're golden.

Installing NNT

Part of inventorado/ComfyUI_NNT - install via ComfyUI Manager (search "ComfyUI Neural Network Toolkit") or clone it into custom_nodes and run pip install -r requirements.txt:

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

That requirements file is a heavy scientific stack (torch, scikit-learn, pandas, shap, seaborn…), so budget a real install and restart ComfyUI afterward. The pack's own example workflows live in its workflows/ folder and want ComfyUI-Jjk-Nodes for the text displays - let Manager install those too. Remember the author's honest framing: NNT is a learning/prototyping tool, not a production training stack - but for learning what normalization does to a training run, it's ideal.

CategoryNNT Neural Network Toolkit/Layers

Inputs (7)

NameTypeDefaultDescription
norm_typeCOMBOBatchNorm2d6 options: None, BatchNorm, LayerNorm, InstanceNorm, GroupNorm, LocalResponseNorm
num_featuresINT641–2048
epsFLOAT0.00001e-10–0.001
momentumFLOAT0.10–1
affineCOMBOTrue2 options: True, False
track_running_statsCOMBOTrue2 options: True, False
LAYER_STACKoptLIST

Outputs (1)

NameTypeDescription
LAYER_STACKLIST