Nodes/ComfyUI-Pt-Wrapper/Ptn LSTM Linear
ComfyUI Node

Ptn LSTM Linear

The LSTM + classifier head combo node for text classification

By HowToSD·Created about a year ago·Updated about a year ago· 7
Ptn LSTM Linear
    • PTMODEL
    input_size1
    hidden_size1
    num_layers1
    biastrue
    batch_firsttrue
    dropout0.000
    bidirectionalfalse
    proj_size0
    linear_output_size1
    linear_biastrue

    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.

    CategoryTraining

    Inputs (10)

    NameTypeDefaultDescription
    input_sizeINT11–1000000
    hidden_sizeINT11–1000000
    num_layersINT11–1000
    biasBOOLEANtrue
    batch_firstBOOLEANtrue
    dropoutFLOAT0.0000–1
    bidirectionalBOOLEANfalse
    proj_sizeINT00–1000000
    linear_output_sizeINT11–1000000
    linear_biasBOOLEANtrue

    Outputs (1)

    NameTypeDescription
    PTMODELPTMODEL