Pt Train Classification Model Lr

Pt Train Classification Model Lr is the pack's workhorse trainer for image classification: it takes a model, a dataloader, an optimizer, and a learning-rate scheduler, runs a full training loop, and hands you back a trained model plus the loss history. The "Lr" in the name is the whole point - unlike the plain training nodes, this one makes a learning-rate scheduler a required input, so it's the version you reach for when you want the rate to decay over training instead of staying flat.

It's part of ComfyUI-Pt-Wrapper, the ~200-node pack that brings PyTorch model building and training into ComfyUI with no Python. This is the node at the heart of the pack's image-classification tutorials - dog-vs-cat on your own images, CIFAR-10 with ResNet (~94% validation accuracy), Fashion MNIST. The loop looks exactly like the PyTorch training loop you'd write by hand, but every piece is a node.

How it works

Under the hood the node runs the standard recipe: for each epoch, iterate the training dataloader, zero the optimizer's gradients, run the model, compute cross-entropy loss (it's baked in - no loss node here), backprop, and step the optimizer. After each epoch it steps the scheduler and, if you've wired a validation loader, evaluates and records validation loss.

Two behaviors worth knowing. First, early stopping: with early_stopping on, the node watches validation loss and halts after early_stopping_rounds epochs without improvement, then restores the best checkpoint. Second, the best model is saved to a temp checkpoint (models/tmp/best_model.pt inside the pack) so output_best_val_model can hand you the best-validation model instead of the last-epoch one - a real quality-of-life detail, because the last epoch is rarely the best one.

The one hard rule here: if you use a ReduceLROnPlateau scheduler (the classic "drop the LR when loss plateaus" scheduler), you must provide a val_loader - the node raises a clear error otherwise, because that scheduler needs validation loss to decide when to lower the rate.

Inputs and outputs

The inputs that matter:

• model (PTMODEL), train_loader (PTDATALOADER), optimizer (PTOPTIMIZER) - the three things every training node needs.
• scheduler (PTLRSCHEDULER) - required here. Build one from the Pto Lr Scheduler family.
• epochs (INT, default 1) - bump this up; 1 epoch does almost nothing.
• use_gpu (BOOLEAN, default false) - flip on for CUDA.
• early_stopping / early_stopping_rounds (10) / output_best_val_model (true) - the quality-of-life trio.
• val_loader (optional) - needed for early stopping, plateau schedulers, and meaningful val loss.

Outputs: Model (PTMODEL) - wire into PtSaveModel or an evaluation node - plus train loss and val loss tensors, the loss-per-epoch history you can route to a text/plot node to watch training.

Installing the pack

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

Restart ComfyUI, or use ComfyUI Manager → search "ComfyUI-Pt-Wrapper" → install (it's on the Comfy Registry; Manager pulls the dependencies - pandas, scikit-learn, transformers and friends). For the best start, drag in an example workflow from examples/workflows/ so the whole data→train→save chain is pre-wired.

Common issues

• Loss never goes down. First check epochs (1 is the default and is basically nothing), then confirm your dataloader actually returns the right shape for the model.
• ReduceLROnPlateau errors without a validation loader. The node tells you this directly - wire in val_loader.
• Best model vs last model confusion. If output_best_val_model is on but you didn't provide val_loader, there's nothing to judge "best" by - keep validation data in the loop if you're relying on it.

This is the node you'll actually reach for on image work. Everything about it is shaped like "train something, don't blow up, hand me the good model."