Pt Train Fine Tune Classification Transformer Model

Pt Train Fine Tune Classification Transformer Model fine-tunes a pretrained Hugging Face transformer for text classification - BERT, RoBERTa, DistilBERT, ALBERT. Instead of building a model from scratch like the other transformer trainer, this one starts from weights that already understand language, then adapts them to your task. With the pack's example workflow, fine-tuning DistilBERT on IMDB sentiment gets you around 92-93% validation accuracy - a big jump over the ~85% you get training a Transformer from zero.

It's part of ComfyUI-Pt-Wrapper, the ~200-node pack that brings PyTorch training into ComfyUI's graph. Where PtTrainClassificationTransformerModel trains your own hand-built model, this node is the "use a real pretrained model" route: you load a Hugging Face model with a linear classification head (Ptn Hf Fine Tuned Classification Model), and this node adapts it.

How it works

The interesting mechanism is the freeze-then-unfreeze schedule, controlled by freeze_pretrained_module_epochs. Training runs like this:

• For the first freeze_pretrained_module_epochs epochs, the node freezes the pretrained module (set_requires_grad(llm_model, False)) - only the new classification head trains, so the pretrained weights don't get wrecked by the early, noisy gradients.
• From freeze_pretrained_module_epochs onward, it unfreezes the whole model and full fine-tuning kicks in.

The pack's own example: freeze_pretrained_module_epochs = 1, epochs = 4 → head-only training for epoch 1, full fine-tune for epochs 2-4. That two-phase schedule is the standard reason fine-tuning works without destroying what the model already knows, and it's a real detail you'd otherwise have to hand-code.

Two constraints worth flagging. The node's docs state that only Adam and AdamW optimizers are supported - pick one from the Pto optimizer nodes. And you provide the loss_function yourself (the text workflows use BCE with Logits Loss for binary sentiment). classification_metrics prints validation accuracy as it trains, and the usual family features - use_gpu, early stopping with early_stopping_rounds, output_best_val_model - all apply.

Inputs and outputs

• model, train_loader, optimizer, loss_function - the core. Optimizer must be Adam/AdamW.
• epochs (INT, default 2) and freeze_pretrained_module_epochs (INT, default 1) - the freeze/unfreeze schedule above.
• use_gpu, early_stopping, early_stopping_rounds (10), output_best_val_model (true), classification_metrics (true).
• val_loader (optional) - needed for early stopping and accuracy printing.

Outputs: Model (PTMODEL), train loss and val loss tensors.

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. The requirements include transformers, so the Hugging Face models download on first use - no manual model setup, but do expect a download the first time you run. Start from examples/workflows/distilbert_classification.json (straight fine-tune) or distilbert_classification_2.json (the freeze/unfreeze schedule); a local-dataset variant (distilbert_classification_local_dataset.json) shows how to swap in your own JSONL data instead of the Hugging Face IMDB set.

Common issues

• Pretrained model won't download. First run needs network access to Hugging Face. If it fails, check your connection and the transformers version before anything else.
• Optimizer not Adam/AdamW. The node is explicit that only these two are supported. Reach for Pto AdamW and you're set.
• Head training destroying the model. If you skip the freeze phase (set freeze_pretrained_module_epochs too low) the early gradient noise can damage pretrained weights - keep at least 1 freeze epoch, and match the model's tokenizer to its name (Hf Tokenizer Encode with the same model name).

It's the highest-accuracy text path in the pack, and the freeze schedule is the detail that makes it work. If you've ever wanted to fine-tune a real BERT-family model without writing a training script, this is the node.