NNT Train Model

This is the heart of the Neural Network Toolkit. Everything else - layer nodes, data loaders, tensor tools - exists to feed this node. NNT Train Model takes a compiled MODEL, your training_data and target_data tensors, runs a real PyTorch training loop, and hands you back a trained model, a training log, and a metrics dictionary you can turn into plots. If you've come to this pack to learn what training a network feels like without writing code, this is where the magic happens.

The mental model that matters: this is not LoRA training or any diffusion fine-tuning you've seen in ComfyUI before. This is a from-scratch, general-purpose neural network trainer - think "the training loop from a Deep Learning course, exposed as a node." Your model comes from NntCompileModel, your data from the loaders, and this node does the rest.

How it works

It wires your tensors into a TensorDataset + DataLoader with your batch_size, moves everything to CUDA if available (CPU otherwise), builds the optimizer from your pick (12 options, default Adam) and the loss from 21 options (default CrossEntropyLoss), then loops for epochs. It handles the details that would normally be boilerplate:

• reshape_input / input_reshape_dim / flatten_input - flatten or reshape your data before training. Both are off by default, which is the trap: a dense network fed raw (32, 1, 28, 28) MNIST batches fails on the first forward pass. Turn on flatten_input (to get (32, 784)) or set reshape_input with a dim like [-1, 784], and a plain dense net suddenly works on MNIST. This is the single most common shape error in the pack.
• Loss-aware target prep - targets get squeezed for CrossEntropy, kept float for MSE. It tries to do the right thing per loss.
• Scheduler - StepLR, ReduceLROnPlateau, or CosineAnnealingLR, with min_lr guarding the floor.
• Early stopping - patience and min_delta, on by default so you don't waste epochs once loss plateaus.
• class_weights - pass a string like [0.5, 1.0] and it becomes a weight= tensor on supported losses, useful for unbalanced classes.

The optional hyperparameters input takes a DICT from the NntTrainingHyperparameters node and overrides the built-in settings - which is your way to keep configs tidy and run experiments without touching twenty fields.

Outputs: model (the trained model, ready to save or run inference), training_log (a string log), and metrics - a DICT with per-epoch loss, accuracy, learning_rates, batch_losses, epoch_times, and best-loss/best-accuracy. That dict is exactly what NntVisualizeTrainingMetrics eats.

What bites

Training is slow, and the pack says so - it's for learning and prototyping, not production. Watch your shapes: if the compile-time input_shape doesn't match what your data actually is, you'll get a runtime shape error at the first forward pass; the flatten/reshape inputs are your lever. And don't expect a 60,000-sample MNIST run to be quick on CPU - slice your data (samples_to_return) down for first experiments. The training log and metrics output are your best friends when loss does something weird: read the log, check the plot.

Install

Pack-level install:

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

Restart ComfyUI (or Manager → "ComfyUI Neural Network Toolkit NNT"). Needs torch and sklearn (both in requirements). The full requirements list is heavy - torch, numpy, pandas, matplotlib, transformers, statsmodels, onnx, shap pinned 0.41.0 - so expect a slow first install. Then build a layer stack, compile it, load a dataset, and let this node do what it does: learn.