ML Forge

A visual PyTorch pipeline editor. Build, train and run image classification models without writing code.

MLForge v1.0.3

Additions:

• Fixed missing package prefix on import in _build_criterion_and_optimizer
• Preserved plot data in last_run on training completion so Metrics window works post-run
• Implemented update_block_labels/reset_block_labels stubs with live epoch stats
• Fixed GaussianBlur sigma range string parsing crashing at runtime
• Derived _DATASET/_AUG/_LOSS/_OPTIMIZER_BLOCKS from blocks.SECTIONS

What it does

• Visual pipeline - drag nodes onto a canvas, connect them with wires, and ML Forge generates and runs the training code for you
• Three-tab workflow - Data Prep -> Model -> Training
• Live training - watch loss curves update in real time, save checkpoints, run inference on your trained model
• Export - export projects into clean PyTorch

Tutorial

Watch the MLForge tutorial video here: MLForge Tutorial

Requirements

IMPORTANT: PyTorch must be preinstalled for training, it is not installed as a dependency.

• Python 3.10 or newer
• PyTorch 2.0 or newer
• torchvision

pip install torch torchvision

GPU training is automatic if CUDA is available. CPU and Apple MPS are also supported.

Quick Start

To install MLForge, enter the following in your command prompt

pip install zaina-ml-forge

Then

ml-forge

Cloning

git clone https://github.com/zaina-ml/ml-forge
python -m ml_forge

Building your first model

1. Data Prep tab

• Add a Dataset node (MNIST, CIFAR10, CIFAR100, FashionMNIST, or ImageFolder)
• Chain transforms: ToTensor is required, add Normalize for best results
• End with a DataLoader (train) node
• For proper validation, add a second chain (same dataset with train=False) ending with DataLoader (val)

2. Model tab

• Start with an Input node - shape is auto-filled from your dataset
• Add layers: Linear, Conv2D, ReLU, BatchNorm2D, Flatten, Dropout, etc.
• End with an Output node - num classes is auto-filled from your dataset
• Connect nodes by dragging from an output pin to an input pin
• in_features and in_channels auto-fill when you connect layers
• After a Flatten node, the next Linear's in_features is calculated automatically

3. Training tab

Add these four nodes from the palette and wire them up:

DataLoaderBlock.images  ->  ModelBlock.images
ModelBlock.predictions  ->  Loss.pred
DataLoaderBlock.labels  ->  Loss.target
Loss.loss               ->  Optimizer.params

Configure epochs, device, checkpointing and early stopping in the right panel, then press RUN.

Keyboard shortcuts

Supported datasets

Inference

After training, open Run -> Inference, browse to your checkpoint (.pth), and click Run Inference to sample from the test set and see top-k predictions.

Metrics

Click the METRICS button to see a summary of your training run: final loss, best validation accuracy, fit diagnosis, and loss/accuracy curves, you may also see the curves on the right training panel.

Saving and loading

Projects are saved as .mlf files (JSON). Use File -> Save / Save As or Ctrl+S.

Exporting code

File -> Export -> Python -> PyTorch generates a standalone train.py that reproduces your pipeline. No ML Forge required to run it.

Future Plans

• Orient to a more beginner-friendly approach
• Add custom block functionality

License

MIT