Running fast with MosaicML Composer and Weight and Biases

MosaicML Composer is a library for training neural networks better, faster, and cheaper. It contains many state-of-the-art methods for accelerating neural network training and improving generalization, along with an optional Trainer API that makes composing many different enhancements easy.

Coupled with Weights & Biases integration, you can quickly train and monitor models for full traceability and reproducibility with only 2 extra lines of code:

from composer import Trainer
from composer.loggers import WandBLogger

wandb_logger = WandBLogger(init_params=init_params)
trainer = Trainer(..., logger=wandb_logger)

W&B integration with Composer can automatically:

• log your configuration parameters
• log your losses and metrics
• log gradients and parameter distributions
• log your model
• keep track of your code
• log your system metrics (GPU, CPU, memory, temperature, etc)

🛠️ Installation and set-up

We need to install the following libraries:

• mosaicml-composer to set up and train our models
• wandb to instrument our training

Use the Composer Trainer class with Weights and Biases 🏋️‍♀️

W&B integration with MosaicML-Composer is built into the Trainer and can be configured to add extra functionalities through WandBLogger:

• logging of Artifacts: Use log_artifacts=True to log model checkpoints as wandb.Artifacts. You can setup how often by passing an int value to log_artifacts_every_n_batches (default = 100)
• you can also pass any parameter that you would pass to wandb.init in init_params as a dictionary. For example, you could pass init_params = {"project":"try_mosaicml", "name":"benchmark", "entity":"user_name"}.

For more details refer to Logger documentation and Wandb docs

let's grab a copy of MNIST from torchvision

we can import a simple ConvNet model to try

📊 Tracking the experiment

we define the wandb.init params here

we are able to tweak what are we logging using Callbacks into the Trainer class.

we include callbacks that measure the model throughput (and the learning rate) and logs them to Weights & Biases. Callbacks control what is being logged, whereas loggers specify where the information is being saved. For more information on loggers, see Logging.

once we are ready to train we call fit

We close the Trainer to properly finish all callbacks and loggers

⚙️ Advanced: Using callbacks to log sample predictions

Composer is extensible through its callback system.

We create a custom callback to automatically log sample predictions during validation.

we add LogPredictions to the other callbacks

Once we're ready to train, we just call the fit method.

We can monitor losses, metrics, gradients, parameters and sample predictions as the model trains.

📚 Resources

• We are excited to showcase this early support of MosaicML-Composer go ahead and try this new state of the art framework.

❓ Questions about W&B

If you have any questions about using W&B to track your model performance and predictions, please reach out to the wandb community.