NectarGAN is comprised of two core components:
A modern, high-performance graphical interface for end-to-end model development. Build datasets, configure architectures, train models, track experiments, and prepare for deployment, all through an intuitive unified UI.
Currently supports:
- Pix2pix-style conditional GANs (paired image translation).
- Full dataset management and augmentation pipeline.
- Real-time training visualization and experiment tracking.
- Model testing and ONNX conversion.
Coming soon:
- Diffusion model training dashboard.
- Classifier training/testing interface.
A fully modular, highly extensible PyTorch-based framework comprised of easy to use building blocks for assembling, training, and testing CV and image synth models. The API offers developers the tools needed to quickly write training/testing scripts with maximum functionality and minimal boilerplate.
from nectargan.trainers.pix2pix_trainer import Pix2pixTrainer
if __name__ == "__main__":
trainer = Pix2pixTrainer(
config='/path/to/config/file.json',
loss_subspec='extended', log_losses=True)
for epoch in range(100):
trainer.train_paired(epoch)
trainer.save_checkpoint()
trainer.save_examples()
trainer.print_end_of_epoch()For more information, please see the NectarGAN API documentation.
- Artists and researchers exploring image synthesis techniques (cGAN, diffusion, autoencoders).
- Startups looking for an easy and flexible way to experiment with and deploy image synth/CV models.
- Students wishing to learn about and explore modern image synthsis techniques in a visual, easy to use environment.
- Engineers/TDs wanting to painlessly integrate computer vision and image synthesis models into their pre-existing pipelines.
Please refer to the quickstart documentation for information on how to get started using NectarGAN.
NectarGAN also supports headless training and testing of models in Docker containers. See below for more information.
| Model | Domain | Currently Available With |
| PatchGAN (Discriminator) | Images (2D) | API · CLI · Toolbox · Docker |
| PatchGAN (Multi-scale) | Images (2D) | API |
| UNet (Generator) | Images (2D) | API · CLI · Toolbox · Docker |
| UNet (Diffusion) | Images (2D) | API |
| ResNet (Generator) | Images (2D) | API |
| ResNet (Classifier) | Images (2D) | API |
| Diffusion (Pixel) | Pixel space | API · CLI |
| Diffusion (Latent) | Latent space | API · CLI |
Toolbox Support: Currently, the Toolbox interface and Docker wrapper support PatchGAN (Discriminator) and UNet (Generator) for Pix2pix-style image translations workflows. All other models are accessible via the API and/or CLI, with Toolbox integration planned for future releases.
Training and inference benchmarks (1x Nvidia RTX 2080Ti, PyTorch 2.9.0 cu128, facades dataset, 256x256):
| Operation | Time | GPU Memory |
|---|---|---|
| Training (per epoch, batch size 1) | ~8.5 seconds | 1.03 GB |
| Training (per epoch, batch size 4) | ~4.0 seconds | 1.18 GB |
| Training (per epoch, batch size 8) | ~3.5 seconds | 1.43 GB |
| Training (per epoch, batch size 16) | ~3.2 seconds | 1.94 GB |
| Inference (single image) | ~30ms | ~0.23 GB |
Diffusion model performance metrics will be added in a future release.
NectarGAN is under active development with regular feature releases and improvements. While the framework is functional and tested, you may occasionally encounter edge cases or bugs. Please report any issues via GitHub so they can be addressed quickly!
Backward Compatibility: Every effort is made to maintain compatibility between releases, but significant feature additions may occasionally require changes to core components. Please see the CHANGELOG for details on any breaking changes.
- Fully modular, PyTorch-based training framework supporting conditional GANs, diffusion models, autoencoders, and classifiers.
- Configurable UNet-style generator with support for multiple UNet block types, and drop-in support for your own custom blocks.
- Modular and highly configurable pixel and latent space diffusion models.
- Easy loss tracking and logging during and after training with a custom LossManager system.
- Modular loss spec system allows you to quickly define reusable objective functions which can be loaded by the LossManager. Also included is a prebuilt Pix2pix objective function with a number of additional subspecs which add some extra loss functions for rapid prototyping.
- The LossManager also supports a custom loss weight scheduling system allowing you to apply one of the included weight schedules, or to easily drop your own custom scheduling function in and have the LossManager take care of the rest.
- Real time inference and loss visualization during training with Visdom (for headless training), or with the NectarGAN Toolbox GUI.
- Hook-based Trainer class allows you to quickly define a custom training loop for your models.
- Easy framework for loading pre-trained model weights to experiment with new training schedules from a earlier checkpoint, or to fine tune a previously trained model.
- Training and testing settings are all handled by JSON configuration files.
- Config copies are automatically exported and version tagged for each train allowing for easy experiment tracking and reproducibility.
- Toolbox UI settings can be loaded in one click from any previously saved config file allowing you to easily pick up training with your previous settings.
- Everything related to a training session is saved in the config file ensuring that no experiment data is lost.
- Automatic config parser allows you to easily add entries to the config file so that you can expand it to suit your model's needs.
- Modern, fully graphical model training, testing, and evaluation tool.
- Configure your model architecture, build your dataset processing chain, define your objective function, and run your training, all with just a few clicks. No code needed, just drag some sliders and see how it changes the models behavior.
- Real time progress and timing stats, inference visualization, and loss graphing with a highly configurable training interface allowing you to put the most focus on whatever metrics you're most interested in seeing.
- A testing interface allowing you to quickly load and validate your models on test datasets. The results will be evaluated with various loss functions and all the images and information will be displayed to you in the interface where you can sort by these metrics.
- Threaded training and testing framework ensures that the UI remains responsive during training and that data visualization has a minimal impact on training and testing speed.
- Training can easily be paused and resumed at any time.
- Load and visualize your model's training example images and loss log data with the review panel. All [x, y, y_fake] example sets will be automatically loaded, and each loss in the log will be graphed with a configurable sample rate.S
- View time statistics during training including slowest, fastest, and average epoch and iteration time, total train time, and a real time time graph of all previous epoch times.
NOTE: Diffusion models will be added to the NectarGAN Toolbox in a future release.
|
A simple but powerful augmentation UI allows you to quickly apply a variety of Albumentations-based augmentations to your datasets at training-time to help expand small datasets and improve model generalization. Anything from random flipping and rotation, to optical distortion and random grayscale can be applied by just dragging a slider. |
See the Docker quickstart guide for more information about using NectarGAN with Docker! |
|
- Pair (A | B) images into Pix2pix input data, extremely quickly, with control over direction and optional image scaling.
- Sort image files by various metrics (white pixel count, black pixel count, mean pixel value, various types of contrast), and then two related tools. One to unsort the files back to their original order, and one to copy to sorting order of one directory of files to another. These can be used to help process and find bad images in large datasets.
- Automatically build train/test/val splits for dataset images with easy control over split percentages.
Easily convert you models to ONNX from either the toolbox UI, or from your own pipeline via the ONNXConverter, and test the resulting model immediately on real images with a live dashboard to display the results.
Please see the FAQs page.
Please see the contribution guide for more information.
- Multi-GPU support.
- More CLI support.
- Exposed normalization options in the Toolbox interface.
- More block types, standard loss functions, scheduling functions, and LossManager specifications.
- Live (and maybe interactive) architecture diagram.
- ONNX vs. PyTorch inference comparison.
- Loss monitoring and early stopping.
- Dataset augmentations preview.
- Houdini tooling. I'd like to build an HDA that lets you run NectarGAN training steps as PDG work items.
This project is inspired by Jun-Yan Zhu's PyTorch-CycleGAN-and-pix2pix implementation.
Please see REFERENCES.md for a complete list of citations.