Nodes/EternalKernel PyTorch Nodes/ExtractLayersAsModel
ComfyUI Node

ExtractLayersAsModel

Slice a sub-network out of your model, frozen or not

By TashaSkyUp·Created about a year ago·Updated about a year ago· 1
ExtractLayersAsModel
  • model
  • model
  • shapes_str
start_idx
end_idx
freezeFalse

The classic transfer-learning move - take a model you built, keep the first few layers as a frozen feature extractor, and retrain only the head - needs a way to slice a model. ExtractLayersAsModel is that: it pulls a contiguous range of layers out of your nn.Sequential and returns them as a brand-new standalone nn.Sequential model. Optionally it freezes every parameter in the slice, which is exactly the "backbone" pattern.

You reach for it when you want one part of a model to behave differently from the rest. Build a feature extractor, ExtractLayersAsModel it out (with freeze=True), then use AddModelAsLayer to graft a fresh trainable head onto it, and train. The frozen extractor stays put while the new head learns - a miniature version of how every fine-tuning workflow on the image side works.

How it works

It checks your model is nn.Sequential, validates the indices, then builds nn.Sequential(*list(model.children())[start_idx:end_idx+1]) - note end_idx is inclusive. It also walks the sliced layers collecting their in_features/out_features into a string (the shapes_str output) so you can see the dimensions of what you extracted. If freeze is True, it sets requires_grad=False on every parameter in the slice. The original model is untouched - you get a new, independent Sequential.

Inputs and outputs

  • model (TORCH_MODEL) - must be a Sequential. Start from SequentialModelProvider and layer-builder nodes.
  • start_idx / end_idx - the inclusive layer range. Layer 0 is the first layer you added. If you don't know the indices, count your builder nodes - or wire shapes_str into a text display to audit.
  • freeze - False copies the layers out trainable; True sets requires_grad=False on everything in the slice.

Outputs:

  • model (TORCH_MODEL) - the extracted sub-model.
  • shapes_str (STRING) - the collected feature sizes as text. Mostly useful for a sanity check or to feed a text node.

Install

ComfyUI Manager, search "EternalKernel PyTorch Nodes", or:

cd ComfyUI/custom_nodes
git clone https://github.com/TashaSkyUp/EternalKernelPytorchNodes
cd EternalKernelPytorchNodes
pip install -r requirements.txt

Restart ComfyUI; node under ETK/pytorch. No model files. Requirements are the standard ComfyUI stack plus scipy, scikit-learn, transformers, einops.

Common issues

  • end_idx out of range or before start_idx. The node validates this and raises a clear ValueError - the range must be within the model's layer count, and start_idx <= end_idx.
  • Inclusive end index. If you want layers 0–3, that's start_idx=0, end_idx=3 - not 2. Off-by-one here silently drops a layer.
  • Frozen means frozen. With freeze=True, training the combined model won't update the extracted backbone's weights at all. If your loss stops moving, that's a sign the only thing training is the head - which might be the point, or might be why it's stuck.
  • Shared parameters in the original. The extracted model shares parameter objects with the source model you sliced from. Training the extractor also trains the original's layers. If you want a true independent copy, be aware of this.

No community tutorials for this pack - but the mechanics are pure nn.Sequential slicing, so they're predictable. Pack quirk: it patches ComfyUI's validator to ignore return_type_mismatch, so check your wire types manually when something behaves oddly.

CategoryETK/pytorch

Inputs (4)

NameTypeDefaultDescription
modelTORCH_MODEL
start_idxINT
end_idxINT
freezeCOMBOFalse2 options: True, False

Outputs (2)

NameTypeDescription
modelTORCH_MODEL
shapes_strSTRING