Pt Train Classification Model

If you want to train your first model in ComfyUI-Pt-Wrapper, this is the node to start with. Pt Train Classification Model is the pack's simplest trainer: it takes a model, a training data loader, and an optimizer, and runs a full classification training loop with cross-entropy loss baked in. No loss node to wire, no scheduler to configure - it's the least intimidating path from "I have a dataset" to "I have a trained classifier."

It's the same machinery that powers the pack's flagship no-code tutorials - the dog-vs-cat image classifier and the FashionMNIST examples in the README run on trainers exactly like this one. The pipeline looks like: PtvDataset (download MNIST/FashionMNIST, or your own image folder) → Pt Data Loader → a model node (linear, conv, ResNet) → an optimizer node (Adam) → this node → done.

The inputs that matter:

• model, train_loader, optimizer - the three required pieces of any training run.
• epochs - passes over the data, default 1 (bump it; one epoch is a smoke test).
• use_gpu - default off. Flip it if you have a GPU, or accept training on CPU.
• early_stopping + early_stopping_rounds - with a val_loader connected, stops training if validation loss stagnates for N rounds.
• output_best_val_model - default true: with a val_loader, it returns the best-validation model instead of the last epoch's.
• val_loader (optional) - a holdout set to measure validation loss each epoch.

Outputs are Model (the trained weights, ready to feed a prediction/evaluation node), train loss, and val loss - per-epoch history tensors you can plot to see learning happen. A val_loader isn't strictly required, but without one, early stopping and best-model selection have nothing to work on, so connect one if you want those features to mean anything.

One behavior to internalize: like all the pack's training nodes, this one re-runs every time you queue - IS_CHANGED always returns NaN. Training with identical inputs is still training. That's expected, just don't leave it wired into a workflow that fires repeatedly.

Install

ComfyUI Manager (search "ComfyUI-Pt-Wrapper"), or:

cd ComfyUI/custom_nodes
git clone https://github.com/HowToSD/ComfyUI-Pt-Wrapper

Restart ComfyUI. Heavy requirements (transformers, datasets, peft, accelerate, scikit-learn, pandas, seaborn, matplotlib, gensim, sentencepiece). No model files; datasets download at runtime.

Common issues

CPU-bound because use_gpu was left off - the classic. Labels that aren't int64 (cast them). And if you see loss stuck at the same value every epoch, check that your model's output dimension matches your number of classes - a mismatch silently produces garbage that cross-entropy can't fix. The best model checkpoint, when saved, lands in the pack's models/tmp/best_model.pt.