Nodes/ComfyDL/Linear Regression Train
ComfyUI Node

Linear Regression Train

A real SGD loop in a node, and it opens its own gradient scope to work at all

By Cynthia-lxx·Created 3 months ago·Updated 13 days ago· 7
Linear Regression Train
  • X
  • y
  • dataset
  • w
  • b
  • loss_history
  • y_hat
◄num_steps100►
◄batch_size32►
◄lr0.030►
◄seed0►

What it is

This is the node that makes ComfyDL make sense. Cynthia-lxx's pack is a from-scratch deep-learning playground built on the d2l textbook code, and Linear Regression Train is the honest entry point: a full mini-batch SGD loop for y = X @ w + b, running inside your ComfyUI graph, with the loss curve coming out of a socket.

It's a teaching node and a plumbing node at once: the d2l linear-regression chapter wired up in ten nodes instead of forty lines of Python. Paste the book's toy dataset into Read String → Dataset, set target to y, and watch it fit. And because w, b and loss_history are ordinary tensors, you can plot them, compare them, or feed them back into other nodes to verify the math.

It's the only training node in the pack's d2l/Training category, which tells you what the pack is going for: the graph is the training script.

How it works

The mechanism is textbook and the source doesn't hide it. Weights and bias start at zeros - w of shape (n_features, 1), b of shape (1,) - and each of num_steps iterations:

  1. draws a batch of row indices with torch.randint(0, n, (batch_size,)), seeded from the seed widget,
  2. computes y_hat = X_batch @ w + b,
  3. takes the mean squared loss with the /2 convention (mean((y_hat - y) ** 2 / 2)),
  4. calls loss.backward(), then does a plain in-place SGD step under torch.no_grad() and zeroes the gradients.

Two implementation details are worth more than the rest of the article.

It opens its own autograd scope. ComfyUI evaluates nodes inside torch.inference_mode(), and autograd cannot cross that boundary - a naive training node dies with "element 0 of tensors does not require grad". The fix in the source is explicit: the whole loop runs inside with torch.inference_mode(False):, and the returned tensors are .detach()ed so nothing leaks a graph back into the workflow. If you've ever wondered why nobody just writes a training node, this is the answer.

Sampling is with replacement, not per-epoch. torch.randint draws batch_size indices independently - it's the stochastic-sampling formulation from the book, not "walk the shuffled set one epoch at a time". Duplicates inside a batch are normal, and batch_size = 0 means full-batch descent, which is the deterministic setting if you want it.

Inputs and outputs

Required: X (TENSOR, [n, f]) and y (TENSOR, [n] or [n, 1]), plus four widgets - num_steps (default 100, up to 100,000), batch_size (default 32; 0 = whole dataset), lr (default 0.03, range 1e-5 to 1) and seed (default 0, which only controls the batch sampling - the initialisation is fixed at zeros, so runs are reproducible).

Optional: dataset (DATASET). Wire a labelled dataset and it wins over X/y. Wrinkle: X and y are declared required in the schema even though the code ignores them when a dataset is present, so if ComfyUI refuses to queue because those sockets are empty, feed them from Dataset → Tensors.

Outputs: w ((n_features, 1)), b ((1,)), loss_history (a 1-D tensor with one value per step) and y_hat (predictions over the training features). loss_history into a Plot gives you the classic descending curve; loss_history[-1] is your sanity metric. w and b wired into a Linear Regression node should reproduce y_hat exactly - that's the verification loop the pack's own docstring describes.

Install

Manager → search ComfyDL, or:

cd ComfyUI/custom_nodes
git clone https://github.com/Cynthia-lxx/ComfyDL
pip install -r ./ComfyDL/requirements.txt

Restart and look for [ComfyDL] Loaded <n> deep learning nodes. in the console. Requirements are just matplotlib, IPython and matplotlib-inline; torch is already there because it's ComfyUI, and pandas comes from your environment. The pack has essentially no community footprint yet, and the author archived the example workflows when a refactor made them stale - so read the node tooltips and the repo's FUNCTIONS.md rather than hunting for a tutorial.

Where people get burned

Learning rate vs feature scale. lr = 0.03 is tuned for textbook data with roughly unit-scale features. Feed it a column in the thousands and the loss will diverge to inf/nan within a few steps. Standardise first (subtract mean, divide by std) - this node does not, unlike the pack's Regression Train node, whose docs mention a train/validation split and feature standardisation. Small lr plus more num_steps is the manual fix if you'd rather not preprocess.

Too few steps. 100 steps is a demo. If the loss curve is still falling when it stops, raise num_steps - the max is 100,000.

A NaN in your data poisons everything. Nothing in the reader→tensor chain drops NaNs, and a single NaN turns the loss into NaN permanently. Clean the file.

Wrong label shape. y is reshaped to (n, 1) for you, and mismatched sample counts raise a clear error naming the counts - the nicest failure message in the pack.

Categoryd2l/Training

Inputs (7)

NameTypeDefaultDescription
XTENSOR—
yTENSOR—
num_stepsINT1001–100000—
batch_sizeINT320–65536—
lrFLOAT0.0300.00001–1—
seedINT00–99999—
datasetoptDATASET—

Outputs (4)

NameTypeDescription
wTENSOR—
bTENSOR—
loss_historyTENSOR—
y_hatTENSOR—