Ptn LSTM Linear
The LSTM + classifier head combo node for text classification
- PTMODEL
If PtnLSTM gives you the raw recurrent layer, this is the version that actually finishes the job for classification. PtnLSTMLinear stacks an LSTM and a linear head into one PTMODEL, and - this is the useful bit - it collapses the LSTM's sequence output by averaging over valid tokens before the head sees anything. You wire in embeddings, you get logits out, no manual pooling step in between. It's the node behind the pack's LSTM text-classification workflow, which gets roughly 87% validation accuracy on IMDB.
How it works
Internally it's an nn.LSTM followed by nn.Linear. The clever part is how it gets from a sequence to a single vector per example: it builds a validity mask from the input itself (a token counts as valid if its row norm is nonzero), averages the LSTM output over those valid tokens, then runs that through the linear layer. That mean-over-valid-tokens behavior is exactly what makes it robust to padded sequences - padding tokens don't drag the average down. If you make the LSTM bidirectional, the head's input width doubles (2 * hidden_size) to match, and the linear weights get Xavier-initialized.
The inputs that matter
All of the LSTM controls plus two:
- input_size - features per token; must match your embedding dimension (e.g. 100 for GloVe in the example workflow).
- hidden_size - the LSTM's hidden dimension.
- num_layers (default
1) - number of stacked LSTM layers. - batch_first (default
True) -(batch, seq, features)layout; note the default differs from raw PyTorch. - dropout - only active between stacked layers.
- bidirectional (default
False) - doubles the hidden dimension feeding the head. - linear_output_size - the number of classes (or 1 for binary). This is what you set to match your labels.
- linear_bias (default
True) - bias on the final layer.
One PTMODEL comes out, which drops straight into Pt Train RNN Model.
Wiring it up
The pack's lstm_classification.json is the reference: Ptv Hf Glove Dataset produces the token embeddings, Pt Data Loader batches them, and PtnLSTMLinear feeds Pt Train RNN Model alongside Pto Adam and Ptn BCE With Logits Loss. The RNN trainer has a use_valid_token_mean flag - when it's on, it averages the outputs the same way this node does, and that consistency is part of why the workflow trains so cleanly. If your loss refuses to go down, that flag is the first thing to check.
Installing
Standard for this pack. ComfyUI Manager → search "Pt-Wrapper", or:
cd ComfyUI/custom_nodes
git clone https://github.com/HowToSD/ComfyUI-Pt-Wrapper
Restart. The pack's dependency list (transformers, datasets, scikit-learn, pandas, etc.) makes the first boot slow, so don't panic at a long install.
Where people get burned
The classic mistake is input_size not matching your embedding dimension - the node won't tell you it's wrong until shapes collide at runtime, and the error can be cryptic. Also, if you're comparing against a raw-PyTorch LSTM you built yourself, remember batch_first=True here. And since IS_CHANGED returns NaN, the model rebuilds on every queue - fine for training, just don't be surprised if evaluation re-runs it too.
Inputs (10)
| Name | Type | Default | Description |
|---|---|---|---|
| input_size | INT | 11–1000000 | — |
| hidden_size | INT | 11–1000000 | — |
| num_layers | INT | 11–1000 | — |
| bias | BOOLEAN | true | — |
| batch_first | BOOLEAN | true | — |
| dropout | FLOAT | 0.0000–1 | — |
| bidirectional | BOOLEAN | false | — |
| proj_size | INT | 00–1000000 | — |
| linear_output_size | INT | 11–1000000 | — |
| linear_bias | BOOLEAN | true | — |
Outputs (1)
| Name | Type | Description |
|---|---|---|
| PTMODEL | PTMODEL | — |