diff --git a/README.md b/README.md index eb1fc73..653f7d5 100644 --- a/README.md +++ b/README.md @@ -1,20 +1,17 @@ ![Logo](docs/images/main_logo.png) -## Update — March 31, 2025 +## Update — April 6, 2025 -**Now with Hugging Face pretrained model support for text classification** +## Now You Can Try Building a Transformer Model from Scratch -You can use **ComfyUI-Pt-Wrapper** to fine-tune Hugging Face's pretrain models for text classification, such as sentiment analysis on a movie review dataset. +You can use **ComfyUI-Pt-Wrapper** to build a Transformer encoder model from scratch for text classification. -You can choose from the following models: -* bert-base-uncased -* roberta-base -* distilbert-base-uncased -* albert-base-v2 +Learn how to construct a Transformer encoder using basic components such as multi-head attention, layer normalization, linear layers, embedding layers, and residual connections. This workflow allows you to train the model for IMDB text classification and achieve around 85% accuracy! -Performance was not compromised in adding this feature — you should achieve around 92-93% validation accuracy on the IMDB dataset with the pretrained DistilBERT model. +![Overview](docs/images/mha_overview.png) +*The complete setup of all required nodes is included in the example workflow.* -Check out [Fine-tuning Hugging Face models for classification](docs/fine_tuning_hugging_face_models_for_classification.md). +Check out the [Building Transformer From Scratch](docs/building_transformer_from_scratch) guide. --- diff --git a/docs/building_transformer_from_scratch.md b/docs/building_transformer_from_scratch.md new file mode 100644 index 0000000..a7d74a1 --- /dev/null +++ b/docs/building_transformer_from_scratch.md @@ -0,0 +1,169 @@ +# Building Transformer From Scratch + +You can use **ComfyUI-Pt-Wrapper** to build a Transformer encoder model from scratch for text classification. + +In this workflow, we demonstrate how to construct a Transformer encoder using basic components such as multi-head attention, layer normalization, linear layers, embedding layers, and residual connections. +Note: This guide does *not* cover implementing multi-head attention from scratch. If you'd like that topic addressed, please open a GitHub issue. + +You can find the complete workflow at: +`examples/workflows/mha.json` + +![Overview](images/mha_overview.png) + +The workflow consists of the following node groups: +* Loading data (green, bottom-left) +* Setting up the model, loss, optimizer, and scheduler (red, bottom-middle) +* Saving a model (blue, bottom-right) +* Defining a custom Transformer model (top, black) + +--- + +## Loading Data + +![Data load](images/mha_dataload.png) + +Use **Ptv Hf Dataset With Token Encode** to load a Hugging Face text dataset. This node uses a SentencePiece tokenizer to convert text into token IDs. + +SentencePiece requires a pretrained model file. A compatible model is included in this package, sourced from [google-t5/t5-base](https://huggingface.co/google-t5/t5-base/tree/main), licensed under Apache 2.0. + +The tokenizer's vocabulary size is **32000**, so this must be specified in the embedding layer. + +**Ptv Hf Dataset With Token Encode**: +- Pads each sentence +- Generates an attention mask (1 for valid tokens, 0 for padding) + +Use **Pt Data Loader** to batch the tokenized dataset for training. + +--- + +## Setting Up Trainer + +![Training](images/mha_trainer.png) + +This workflow uses: +- **Pto AdamW** as the optimizer +- **Pto Lr Scheduler Cosine Annealing** as the learning rate scheduler + +You can replace these with other optimizers/schedulers available in the extension. + +Loss function: **Ptn BCE with Logits Loss** +This includes a built-in sigmoid activation, so raw logits can be passed directly. + +For training: +- Use **Pt Train Classification Transformer Model** +- Set `unpack_input = True` + This is required because the Data Loader outputs a tuple (tokens, attention mask), but the model expects unpacked inputs. + +If the loss doesn’t decrease, check that this flag is correctly set. + +You can enable `classification_metrics` to print validation accuracy during training. + +--- + +## Model Saving + +Use **Pt Save Model** to save the trained model to disk. + +![Model saving](images/mha_model_save.png) + +--- + +## Model + +A Transformer encoder is composed of one or more encoder blocks. For instance, OpenCLIP (used in Stable Diffusion) has 12. + +This example builds a model with a *single* encoder block: + +![MHA model](images/mha_model.png) + +Breakdown: +* **Embedding block** (top orange) +* **Attention block** (middle green) +* **Feedforward block** (bottom blue) +* **Classification head** (not shown) + +--- + +## Embedding Block + +![Embedding block](images/mha_embedding.png) + +The embedding layer maps token IDs to trainable vectors. + +* num_embeddings: Specify the vocabulary size of tokens. +* embedding_dim: Specify the dimension of each word embedding. In this case 256 is specified in a separate constant node and connected to this node. You can hard code the number, but it's easier to maintain if you define the embedding dimension in one place. + +For example, with a vocabulary of 32,000 and an embedding size of 256, the embedding matrix is [32000 × 256]. + +--- + +## Attention Block + +![Attention block](images/mha_attention.png) + +The attention block learns contextual relationships between tokens and transforms the sequence into a richer representation suited for downstream tasks (e.g., classification). + +--- + +## Feedforward Block + +![Feedforward block](images/mha_ff.png) + +This block consists of two linear layers and further refines the output from the attention block to better encode the sentence meaning. + +--- + +## Classification Head + +![Workflow](images/mha_head.png) + +Transformers output 3D data: `(batch, sequence length, feature dim)`. +For classification, this must be reduced to `(batch, prediction)`. + +Common reduction strategies include: +- Using a specific token (e.g., [CLS]) +- Taking the mean across the sequence + +This workflow uses **mean pooling**, but **ignores padding tokens** when averaging. + +The result is passed to a linear layer producing one logit per example. +No sigmoid is applied here, as the loss function node includes it. + +--- + +## Residual Block + +Residual connections are a key innovation in deep learning. They: +1. Save the input +2. Process it +3. Add the original input back to the processed output + +Transformers apply residuals in both the attention and feedforward blocks. +This workflow uses **Ptn Residual** to wrap submodels with a residual connection. + +--- + +## Model Chaining + +![Model chaining](images/mha_chain.png) + +Yellow blocks represent **model chaining**. For instance: +- The Embedding layer feeds into the Attention layer +- The Attention layer feeds into the Feedforward layer + +Use **Ptn Chained Model With Attention Mask** to chain layers that require the attention mask (e.g., Multihead Attention). + +This node handles: +- Passing attention masks only to layers that need them +- Maintaining correct data flow + +You must explicitly specify which layers consume the attention mask. + +--- + +## Expected Accuracy + +With this workflow, you should reach around **85% validation accuracy** on the Hugging Face IMDB dataset. +This is a solid performance for a Transformer built entirely from scratch. + +--- diff --git a/docs/images/mha_attention.png b/docs/images/mha_attention.png new file mode 100644 index 0000000..8374aa9 Binary files /dev/null and b/docs/images/mha_attention.png differ diff --git a/docs/images/mha_chain.png b/docs/images/mha_chain.png new file mode 100644 index 0000000..27d667b Binary files /dev/null and b/docs/images/mha_chain.png differ diff --git a/docs/images/mha_dataload.png b/docs/images/mha_dataload.png new file mode 100644 index 0000000..a9cb3b5 Binary files /dev/null and b/docs/images/mha_dataload.png differ diff --git a/docs/images/mha_embedding.png b/docs/images/mha_embedding.png new file mode 100644 index 0000000..0bd3cd1 Binary files /dev/null and b/docs/images/mha_embedding.png differ diff --git a/docs/images/mha_ff.png b/docs/images/mha_ff.png new file mode 100644 index 0000000..de4aac7 Binary files /dev/null and b/docs/images/mha_ff.png differ diff --git a/docs/images/mha_head.png b/docs/images/mha_head.png new file mode 100644 index 0000000..208b3af Binary files /dev/null and b/docs/images/mha_head.png differ diff --git a/docs/images/mha_model.png b/docs/images/mha_model.png new file mode 100644 index 0000000..49ca00c Binary files /dev/null and b/docs/images/mha_model.png differ diff --git a/docs/images/mha_model_bottom.png b/docs/images/mha_model_bottom.png new file mode 100644 index 0000000..9e13664 Binary files /dev/null and b/docs/images/mha_model_bottom.png differ diff --git a/docs/images/mha_model_save.png b/docs/images/mha_model_save.png new file mode 100644 index 0000000..c65100b Binary files /dev/null and b/docs/images/mha_model_save.png differ diff --git a/docs/images/mha_overview.png b/docs/images/mha_overview.png new file mode 100644 index 0000000..d828e02 Binary files /dev/null and b/docs/images/mha_overview.png differ diff --git a/docs/images/mha_trainer.png b/docs/images/mha_trainer.png new file mode 100644 index 0000000..1e7b751 Binary files /dev/null and b/docs/images/mha_trainer.png differ diff --git a/pyproject.toml b/pyproject.toml index 0eada01..3c1e6da 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "pt-wrapper" description = "PyTorch extension for ComfyUI featuring extensive PyTorch wrapper nodes for seamless tensor operations and PyTorch model training." -version = "4.0.0" +version = "4.0.1" license = {file = "LICENSE"} [project.urls]