Pt Train Classification Transformer Model

Pt Train Classification Transformer Model trains a Transformer encoder for text classification - and it's the node that makes the pack's flagship "build a Transformer from scratch" workflow possible. The whole point of that workflow is that you construct a multi-head-attention encoder out of basic layer nodes (attention, layer norm, linear, embedding, residual connections), then train it here, no Python anywhere. On IMDB sentiment it lands around 85% validation accuracy, which is respectable for a model you assembled from Lego bricks.

It's part of ComfyUI-Pt-Wrapper, the ~200-node pack that brings PyTorch training into ComfyUI's graph. Where PtTrainClassificationModelLr is the image-workhorse, this is the text-classification trainer, with a couple of extra inputs that matter specifically for tokenized, attention-masked data.

How it works

The training loop is the standard one: iterate the dataloader, run the model, compute the loss you've wired in, backprop, step the optimizer, step the scheduler. Three things make this node different from the simpler trainers:

• loss_function is an input. You pick the loss from the Ptn loss nodes. The pack's text workflows use BCE with Logits Loss - binary classification, logits straight from the model, sigmoid built in.
• unpack_input. This is the one that bites. The dataloader outputs a tuple - (token_ids, attention_mask) - but your Transformer's forward takes them as separate arguments. Set unpack_input = True and the node unpacks the batch before calling the model. The pack's own guide is blunt: if your loss doesn't decrease, check this flag first.
• classification_metrics. With this on (and a val_loader connected), the node prints validation accuracy to the console as it trains, using scikit-learn's accuracy_score. That's your live "is this actually learning?" signal.

Early stopping, best-model tracking, and GPU use all work like the rest of the trainer family: early_stopping + early_stopping_rounds halts on plateauing validation loss, output_best_val_model returns the best-checkpoint model rather than the last epoch's.

Inputs and outputs

• model, train_loader, optimizer, loss_function - the required core.
• epochs (INT, default 1) - raise it; 1 epoch is a smoke test at best.
• use_gpu (default false), early_stopping (false), early_stopping_rounds (10), output_best_val_model (true), classification_metrics (true).
• Optional: scheduler (PTLRSCHEDULER), val_loader, and the critical unpack_input (false).

Outputs: Model (PTMODEL), train loss and val loss tensors (the per-epoch histories, wire them to a display node to watch the curves).

Installing the pack

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

Restart ComfyUI, or ComfyUI Manager → search "ComfyUI-Pt-Wrapper" → install (on the Comfy Registry; dependencies include pandas, scikit-learn, transformers and sentencepiece). Start from the mha.json example workflow in examples/workflows/ - it has the full Transformer, data, and trainer chain pre-assembled, including the SentencePiece tokenizer the pack ships (T5's, vocab size 32000, which is why your embedding layer should say 32000 too).

Common issues

• Loss flat / not decreasing. unpack_input is the first suspect - if the model gets a tuple where it expects arguments, nothing learns.
• Accuracy stuck near random. Check that your embedding dim and vocab size (32000) match the tokenizer, and that the model output dim matches your number of classes.
• Wrong loss for the task. BCE-with-logits for binary, cross-entropy for multi-class - mismatch shows up as loss that doesn't move.

It's the payoff node of the whole pack: everything about Transformers you've read about, rebuilt as nodes, and trained right there on your desktop.