Skip to content

Examples¤

We provide full examples of Solstice usage in different settings. Each example is runnable as a standalone script. You can set up an environment for running the examples in just a few clicks by using the provided devcontainer; see the Contributing page for more information.

mnist_from_scratch ¤

The MNIST example demonstrates how to implement everything yourself using just the Solstice base classes and Haiku to define the neural network.

Summary

This example implements:

  • a solstice.Metrics class for keeping track of accuracy and loss

  • a solstice.Experiment class for specifying how to train the Haiku model

  • basic custom training and evaluation loops for running the experiment

resnet_classification ¤

The ResNet example demonstrates how to perform image classification on CIFAR10 with data-parallel multi-gpu support. Uses Haiku to define the base neural net and Hydra for config management.

Faq

Run python examples/resnet_classification.py -h to see the available configuration options.

Logging is done with TensorBoard. Run tensorboard --logdir=outputs to view the logs.

Summary

This example demonstrates:

  • How to implement solstice.Experiment for training a ResNet50 with multi-gpu support.

  • Usage of solstice.ClassificationMetrics for tracking metrics.

  • Usage of solstice.LoggingCallback with TensorBoard (using the CLU SummaryWriter interface).

  • Usage of solstice.ProfilingCallback for profiling with TensorBoard.

Warning

Multi-GPU support is not yet implemented. Currently only working with single GPU.

adversarial_training ¤

The adversarial training example demonstrates how to implement custom training logic to remove bias from Colour MNIST (based on https://arxiv.org/abs/1812.10352). We use Equinox to define the base network.

Summary

This example demonstrates:

  • How to implement solstice.Experiment for adversarially training a fair classifier.

  • How to implement a custom solstice.Metrics class for tracking fairness-related metrics.

  • How to implement a custom solstice.Callback for conditional early stopping.

  • Usage of solstice.LoggingCallback with Weights and Biases integration.

Warning

Aspirational, not implemented yet.

vmap_ensemble ¤

The vmap ensemble example demonstrates how to implement different implement parallelism strategies by training an ensemble of small neural networks simultaneously on one GPU (inspired by https://willwhitney.com/parallel-training-jax.html). Uses Flax to define the base network.

Summary

This example implements:

Warning

Not Implemented Yet

x_validation ¤

The X validation example demonstrates how to implement K folds cross validation with a custom training loop.

Summary

This example implements:

Warning

Not Implemented Yet