ML Forge: a node editor that generates PyTorch training code for image classifiers

The gap ML Forge fills between a notebook and a GUI trainer

Most GUI training tools hide the code. ML Forge does the opposite: the README states that the graph generates and runs training code, and that File -> Export -> Python -> PyTorch writes a standalone train.py that reproduces your pipeline with no ML Forge required to run it. That export path is the reason to look at this project rather than a generic no-code trainer. The artifact you keep is ordinary PyTorch.

The audience is narrow and the README is honest about it. The topics list includes beginner-friendly and no-code, and the Future Plans section says the project intends to orient toward a more beginner-friendly approach. So the current version is a beginner tool that still expects a working PyTorch environment. The requirements section is blunt: PyTorch must be preinstalled for training, it is not installed as a dependency. If you cannot run pip install torch torchvision yourself, ML Forge will launch and then fail at the point where you press RUN.

Three tabs, two graphs, and one wiring contract

The workflow is split into Data Prep, Model and Training, and the split matters because each tab produces a different kind of object. Data Prep builds a chain: a Dataset node (MNIST, CIFAR10, CIFAR100, FashionMNIST or ImageFolder), then transforms, with ToTensor described as required, ending in a DataLoader (train) node. The README notes that validation needs a second chain using the same dataset with train=False, ending in DataLoader (val). Nothing stops you from skipping that second chain, and nothing warns you either; you simply get a run with no validation curve.

The Model tab is a layer graph. It starts with an Input node whose shape is auto-filled from the dataset, continues through Linear, Conv2D, ReLU, BatchNorm2D, Flatten, Dropout and similar nodes, and ends at an Output node whose class count is auto-filled. Shape inference is the mechanism that makes this usable: in_features and in_channels auto-fill when you connect layers, and after a Flatten node the next Linear's in_features is calculated automatically. That is the part of the editor doing real work rather than just drawing boxes.

The Training tab is a third graph, and the README gives its wiring as a literal block: DataLoaderBlock.images to ModelBlock.images, ModelBlock.predictions to Loss.pred, DataLoaderBlock.labels to Loss.target, and Loss.loss to Optimizer.params. Four nodes, four wires. Epochs, device, checkpointing and early stopping live in the right panel. Because the pin names are fixed strings, a typo in your mental model of the graph produces a disconnected graph rather than an error message you can act on. That is a design choice worth knowing before you start.

Install, launch, and the two supported entry points

The README gives two paths. For the packaged route: pip install zaina-ml-forge, then the command ml-forge. For a source checkout: git clone https://github.com/zaina-ml/ml-forge followed by python -m ml_forge. Note that the PyPI distribution name (zaina-ml-forge) and the repository name (ml-forge) differ, which is a common source of confusion when you search for either one.

Prerequisites are stated as Python 3.10 or newer, PyTorch 2.0 or newer, and torchvision, installed separately with pip install torch torchvision. Device selection is described as automatic: GPU training if CUDA is available, with CPU and Apple MPS also supported. The v1.0.4 release notes mention a patch for a segfault on Mac during template loading, so if you are on macOS and hit that crash, the fix is in the newest release rather than in an older pin.

Projects save as .mlf files, which the README describes as JSON. That means your pipeline is diffable and inspectable in a text editor, which is a practical advantage over binary project formats. Ctrl+S saves, Ctrl+Z and Ctrl+Y undo and redo, Del or Ctrl-Backspace deletes selected nodes, and middle-drag pans the canvas.

Datasets are a closed list, and that is the main constraint

The supported dataset table has five rows: MNIST and FashionMNIST at 10 classes and 1 x 28 x 28, CIFAR-10 at 10 classes and 3 x 32 x 32, CIFAR-100 at 100 classes and 3 x 32 x 32, and ImageFolder with custom classes at 3 x 224 x 224. ImageFolder is the escape hatch for your own data, but it assumes a directory-per-class layout and RGB images, and it is fixed at 224 x 224 in the table. Grayscale, non-square, or multi-label data does not appear to be covered.

This is where the tool stops being general. There are no nodes listed for object detection, segmentation, keypoints, text, audio or tabular input. The layer palette in the README is a CNN vocabulary (Conv2D, BatchNorm2D, Flatten, Dropout, Linear, ReLU). If your problem is not single-label image classification, ML Forge is not a smaller version of the right tool; it is the wrong tool. The README's own Future Plans list only two items, a more beginner-friendly orientation and custom block functionality, so extension beyond the current node set is planned rather than present.

A second constraint is invisible until it bites: transforms are chained by hand, and the README flags ToTensor as required and Normalize as recommended for best results. There is no automatic augmentation policy. What you wire is what runs.

Training, metrics and inference inside the app

Once the four-node training graph is wired, you configure epochs, device, checkpointing and early stopping in the right panel and press RUN. The README describes live training with loss curves updating in real time, plus checkpoint saving and inference on the trained model. The METRICS button opens a summary of the run: final loss, best validation accuracy, a fit diagnosis, and loss and accuracy curves, which also appear in the right training panel.

Inference is a separate flow. After training, you open Run -> Inference, browse to a .pth checkpoint, and click Run Inference, which the README says samples from the test set and shows top-k predictions. Two things are worth flagging. First, the fit diagnosis is described but not defined in the README, so you cannot know from the documentation what rule it uses to call a run overfit or underfit. Second, the inference step samples from the test set, not from arbitrary user-supplied images, so it is a sanity check on the checkpoint rather than a deployment path. If you need to classify a folder of new images, the exported train.py plus your own inference loop is the route the README actually supports.

Checkpointing and early stopping are listed as panel settings rather than described in detail. Treat them as configuration you will verify by running a short job, not as behaviour you can predict from the docs.

How ML Forge differs from a notebook or a scripted trainer

The obvious alternative for this job is a Jupyter notebook with a hand-written PyTorch training loop, or a scripted trainer like the torchvision training references. The difference is not capability. A notebook can express anything ML Forge can, plus detection, segmentation, custom losses and multi-GPU. The difference is where the structure lives. In a notebook, the graph exists only in your head and in the order of cells; in ML Forge it is an explicit, saveable .mlf document with typed pins and automatic shape propagation, which is why the editor can auto-fill in_features after a Flatten node and a notebook cannot.

That trade runs the other way too. A notebook handles a custom loss function in ten lines. ML Forge handles it only if a node exists for it, and the README lists no custom-loss node; custom block functionality is in Future Plans. So the graph is a strong scaffold for the standard classification loop and a wall for anything outside it.

Against a heavier visual framework, the distinguishing choice is the export. The README's claim is that File -> Export -> Python -> PyTorch generates a standalone train.py that reproduces the pipeline and runs without ML Forge. That makes the tool a code generator with a GUI front end rather than a runtime you stay inside. A framework whose graphs only run inside its own runtime asks for a different kind of commitment.

Release cadence, licence and what to check before you commit

The repository is MIT-licensed, not archived, and the recent releases are v1.0.2 in March 2026, v1.0.3 later the same month, and v1.0.4 in August 2026. The v1.0.4 notes list three changes: a patch for a Mac segfault during template loading, a logo added to the splash window, and hover hints added to the node palette for certain nodes. Two of the three are cosmetic, which suggests a project in a polish phase rather than one adding capability. The version numbering has been 1.0.x across all three releases, so no breaking-change signal is available from the tags alone.

Maintenance cost for you is low in one respect and non-trivial in another. Low, because the output is plain PyTorch: if ML Forge stalls, your exported train.py keeps working, and .mlf projects are JSON, so a stalled editor does not lock up your pipelines. Non-trivial, because the README makes PyTorch an external prerequisite rather than a pinned dependency. You own the torch and torchvision versions, which means you also own any incompatibility between them and the nodes ML Forge emits. The README states a floor of PyTorch 2.0 and Python 3.10, not a ceiling, so a future torch release is untested territory until the project says otherwise.

Under MIT you may use, modify and redistribute the code, including in closed products, provided the licence notice is preserved. That is a summary of the licence text, not legal advice; read the LICENSE file in the repository if the distinction matters to you. The practical question before adopting is narrower: does the exported train.py match what you would have written, and does the node palette cover your model? Build one pipeline end to end, export it, and read the generated file. That single check answers both.

Editorial conclusion

Adopt ML Forge if you are teaching image classification, prototyping a small CNN on MNIST, FashionMNIST, CIFAR-10, CIFAR-100 or an ImageFolder tree, or if you want a graph that exports to a standalone train.py you can read. Do not adopt it for segmentation, detection, text, audio, tabular data or multi-GPU work; the README lists no nodes for those, and the dataset list is fixed to five entries. Before committing, run pip install zaina-ml-forge, check that your torch and torchvision versions satisfy the README's floor of Python 3.10 and PyTorch 2.0, build the four-node training graph exactly as the README specifies, and confirm that File -> Export -> Python -> PyTorch produces a train.py that runs without ML Forge installed.

Official sources

1. Issues
2. License: MIT
3. README
4. Releases
5. zaina-ml/ml_forge on GitHub