LibTorch Playground

Blazingly slow flow matching model DDP training/generation implemented with LibTorch

Prerequisites

• vcpkg
• cmake
• ninja
• OpenCV
• CIFAR-100 dataset (i.e. cifar-100-binary/train.bin)
• connection to GitHub
• (optional) astral uv
• (optional) ffmpeg

Usage

Build

At root of this repo:

1. uv venv && . .venv/bin/activate && uv pip install torch (i.e. install LibTorch)
2. cmake --preset linux && cd build/linux
3. ninja nn

Train

1. cd build/linux
2. ./nn config -t -m to generate JSON configuration files
3. Edit the configuration files
4. ./nn train model.json train.json to train from scratch

To train with Distributed-Data-Parallel (DDP), launch multiple instances of the application nn with different --rank, --local-rank, --world-size, --local-world-size commandline options. (See ./nn train -h for details)

Generation

1. cd build/linux
2. ./nn config -e to generate JSON configuration file
3. Edit the configuration file
4. Run ./nn eval model.json DiT_4000.pt eval.json to generate with model weights from DiT_4000.pt
5. The generation results are saved to generated_image.png (and optionally an output.mp4) under CWD

Note: by default, the generation process will be captured as an MP4 output.mp4 by piping every frame to FFmpeg. Compile with template variable Visual = false in torch.cpp to disable.

AI Usage

Documentation, boilerplates and LibTorch internals discovering.