Pt Train Regression Model

Pt Train Regression Model trains a model to predict continuous values - a temperature, a price, a quality score - rather than a category. Where the classification trainers output "dog or cat," this one outputs a number, and it's the node you reach for whenever your target isn't a label but a quantity. If your training data has real-valued targets, this is your trainer.

It's part of ComfyUI-Pt-Wrapper, the ~200-node pack that puts PyTorch model building and training into ComfyUI's graph. It's the regression sibling of PtTrainClassificationModelLr and shares its shape: model in, dataloader in, optimizer in, trained model and loss history out. The difference is underneath, in the loss function and what "good" means.

How it works

Mechanically it's the standard training loop - iterate the dataloader, forward pass, compute loss, backprop, step the optimizer, repeat for each epoch. Two things distinguish it from the classification trainers:

• Mean squared error is baked in. The node uses nn.MSELoss() internally - there's no loss_function input, no scheduler input. That's the natural loss for regression: it punishes big mistakes quadratically, which is usually exactly what you want when predicting continuous values. It also means less to configure: model, data, optimizer, epochs, and you're training.
• No classification metrics. There's no accuracy to print for a regression - the metric is the loss itself. You watch the train loss / val loss outputs instead, and they should fall steadily as the model learns.

Everything else matches the family: use_gpu to train on CUDA, early_stopping + early_stopping_rounds (10) to halt when validation loss plateaus, output_best_val_model (true) to return the best-validation checkpoint rather than the last epoch's, and an optional val_loader for validation. The best-model checkpoint is stashed under models/tmp/ inside the pack, so switching to best-over-last is free.

Inputs and outputs

The ones you'll actually touch:

• model (PTMODEL), train_loader (PTDATALOADER), optimizer (PTOPTIMIZER) - the core three. Your model's output dim should be 1 (a single predicted number) for most regression tasks.
• epochs (INT, default 1) - raise it; 1 epoch is a smoke test.
• use_gpu (default false), early_stopping (false), early_stopping_rounds (10), output_best_val_model (true).
• val_loader (optional) - wire one in if you want early stopping or meaningful validation loss.

Outputs: Model (PTMODEL), plus train loss and val loss tensors - per-epoch loss histories you can route to a display node.

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; dependencies like pandas, scikit-learn, transformers and the rest are pulled automatically). The pack's docs lead with image and text classification examples, so for regression you'll be assembling the flow yourself: a dataset node → dataloader → model → this trainer → PtSaveModel or an evaluation node.

Common issues

• Loss flat from the start. First suspect: your targets aren't scaled. Regression hates huge or wildly varying target magnitudes - normalizing targets (and features) to a sane range is the classic fix, and it's the first thing to check before blaming the node.
• MSE hiding small errors. A mean-squared loss on data with large values can look terrible while being fine in relative terms - plot the loss curves and compare against your target scale.
• Model output dim mismatch. If your model's last layer outputs more than one value, your loss shape won't match your targets. For regression, output dim 1 is the usual setup.

It's the quieter member of the trainer family, but there's no substitute when your task is "predict a number." Fewer knobs than its classification siblings, and that's a feature.