NNT Fine Tune Model

If you came here from the image-generation side of ComfyUI expecting to fine-tune a diffusion checkpoint, stop. NNT Fine Tune Model is not that, and that's actually the point. The Neural Network Toolkit is a teaching sandbox from inventorado - a hobby project for building and training small PyTorch networks visually, node by node. Here "fine-tuning" means taking one of the little nn.Modules you compiled with NntCompileModel and training it further on a new dataset. Think MNIST, not Flux. The tool is explicit about this: it's a learning tool, not a production training rig.

What it does

It's the second half of the training story in this pack. NntTrainModel is the from-scratch trainer; this node is the "keep going" version. You feed it a compiled model plus train/validation splits as tensors, pick a loss, an optimizer, and how many epochs to run, and it returns a fine-tuned model and a text log of per-epoch loss.

The inputs that matter most:

• MODEL - the compiled NNT model (the model output of NntCompileModel).
• train_data / train_labels / val_data / val_labels - four TENSOR inputs. Your dataset split, already turned into tensors by the pack's data-processing nodes. Labels matter: the default loss is CrossEntropyLoss, so plain classification labels in [0, num_classes) are what it expects.
• epochs (default 10) and batch_size (default 32) - the two knobs you'll actually touch.
• loss_function and optimizer - dropdowns over the full PyTorch menu (21 losses, 12 optimizers), defaulting to CrossEntropyLoss and Adam.
• optimizer_params - a JSON-ish string like {"weight_decay": 1e-5} that gets ast.literal_eval'd and passed to the optimizer. Empty {} is fine.
• early_stopping and early_stopping_patience - stops if validation loss stalls for N epochs.
• save_best_model and best_model_path - saves the best state dict to a .pth file (default best_model.pth, relative to your ComfyUI working directory) and reloads it at the end.

Outputs: fine_tuned_model (MODEL) and training_log (STRING - the per-epoch loss lines, handy to pipe into a text-display node).

How it actually runs

Under the hood it's a plain PyTorch loop: TensorDataset + DataLoader, optimizer.zero_grad(), loss.backward(), optimizer.step(), then validation. You can enable a scheduler (StepLR, ReduceLROnPlateau, ExponentialLR, CosineAnnealingLR) via the scheduler dropdown. One quirk worth knowing: the validation pass doesn't bother wrapping itself in torch.no_grad(), so it burns a bit of gradient work it doesn't need. That's the kind of thing that's fine in a teaching tool and would get you a stern look in production code.

Common issues

• "Fine-tuning" confusion - the biggest trap. Readers coming from LoRA-land expect tiny adapters on a big frozen model. NNT Fine Tune Model trains the whole model you hand it. If you want to freeze layers, do it with the edit/merge nodes before this one.
• Label shape mismatches - CrossEntropyLoss wants integer class indices, not one-hot vectors. The pack's NntDatasetToTargetTensor defaults to sparse encoding for exactly this reason.
• It can be slow - a 10-epoch run of a CNN on CPU for CIFAR-sized data takes a while. Lower epochs, or shrink batch_size if your GPU (or lack of one) is struggling.

Install

Search "ComfyUI Neural Network Toolkit NNT" in ComfyUI Manager, or:

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

Standalone installs use the embedded Python: ..\..\..\python_embeded\python.exe -m pip install -r requirements.txt. The requirements list is long (transformers, datasets, onnx, shap, statsmodels...) - the pack pulls in the whole science stack, so first launch after install can take a while. Expect the usual ComfyUI restart after installing.

It's a niche node - 57 impressions in a year tells you this pack lives in a small corner of the ecosystem. But if you're teaching yourself neural networks and want to watch weights actually move, it's the most honest way to do it in ComfyUI.