Visualize Keras Models with One Line of Code

Lukas & Jeff

I love how simple and clear Keras makes it to build neural networks.  With wandb, you can now visualize your networks performance and architecture with a single extra line of python code.

Just add “from wandb import magic” to the top of your training script

To test this functionality, I modified a few scripts in the Keras examples directory.

To install wandb, just run “pip install wandb” and all of my Keras examples should work for you.

1. Simple CNN

I started with the requisite

I added the “from wandb import magic” line below - you can also look at my forked from the Keras examples with the one line change.

Now when the model runs, wandb starts a process in the background saving relevant metrics and streaming them to  You can go to and look at the output of my run.

I can see exactly the data that my model is labeling and view the loss and accuracy curves automatically.

2. Resnet on Cifar

Next, I forked and made the same one line change.  You can see a nice visualization of a resnet at

On the system page, I can see that this model is using a little more of my single GPU than the mnist example.

3. Siamese network

Next I tried the siamese network example.  Here I might want to look at the TensorFlow graph, luckily with our one line of code we automatically instrument and host TensorBoard.  You can find this run at

This instrumentation took me under a minute per model, adds very little compute overhead, and should work for any Keras model you are working on.  As you want to track more things you may want to replace the one line with:

import wandb

Then you can use our custom wandb.log() function to save anything you want.  You can learn more in our documentation.

I really hope you find this useful!

Join our mailing list to get the latest machine learning updates.