Neptune - TensorFlow / Keras Integration

Note

Neptune integrates with both TensorFlow / Keras directly and via TensorBoard. If you are already using TensorBoard you may want to see the TensorBoard integration page

What will you get with this integration?

TensorFlow is an open source deep learning framework commonly used for building neural network models. Keras is an official higher level API on top of TensorFlow. Neptune helps with keeping track of model training metadata.

With Neptune + TensorFlow / Keras integration you can:

  • log hyperparameters for every run

  • see learning curves for losses and metrics during training

  • see hardware consumption and stdout/stderr output during training

  • log TensorFlow tensors as images to see model predictions live

  • log training code and git commit information

  • log model weights

Tip

You can log many other experiment metadata like interactive charts, video, audio and more. See the full list.

Note

This integration is tested with tensorflow==2.3.1, neptune-client==0.4.132, and neptune-contrib==0.25.0

Where to start?

To get started with this integration, follow the quickstart below. You can also skip the basics and take a look at how to log model weights and prediction images in the more options section.

If you want to try things out and focus only on the code you can either:

Quickstart

This quickstart will show you how to:

  • Install the necessary Neptune packages

  • Connect Neptune to your TensorFlow / Keras model training code and create the first experiment

  • Log metrics, training scripts and .git info to Neptune

  • Explore learning curves in the Neptune UI

Note

This quickstart shows only Neptune-specific code snippets. You can see the full code example in Colab/Github or Neptune. Click on one of the buttons at the top or bottom of this page.

Before you start

You have Python 3.x and following libraries installed:

pip install tensorflow==2.3.1 neptune-contrib neptune-client

You also need minimal familiarity with TensorFlow / Keras. Have a look at this TensorFlow tutorial to get started.

Step 1: Initialize Neptune

Add the following snippet at the top of your script.

import neptune

neptune.init(api_token='ANONYMOUS', project_qualified_name='shared/tensorflow-keras-integration')

Tip

You can also use your personal API token. Read more about how to securely set the Neptune API token.

Step 2: Create an experiment

Run the code below to create a Neptune experiment:

neptune.create_experiment('tensorflow-keras-quickstart')

This also creates a link to the experiment. Open the link in a new tab. The charts will currently be empty, but keep the window open. You will be able to see live metrics once logging starts.

When you create an experiment Neptune will look for the .git directory in your project and get the last commit information saved.

Note

If you are using .py scripts for training Neptune will also log your training script automatically.

Step 3: Add NeptuneMonitor Callback to model.fit()

Import NeptuneMonitor() callback from the neptunecontrib package and pass it to the callbacks argument of model.fit().

from neptunecontrib.monitoring.keras import NeptuneMonitor

model.fit(x_train, y_train,
          epochs=42,
          callbacks=[NeptuneMonitor()])

Step 4: Run your training script

Run your script as you normally would:

python train.py

Step 5: Monitor your TensorFlow / Keras training in Neptune

Now you can switch to the Neptune tab which you had opened previously to watch the training live!

TensorFlow Keras learning curve charts

More Options

Log hardware consumption and stderr/stdout

Neptune can automatically log your CPU and GPU consumption during training as well as stderr and stdout from your console. To do that you just need to install psutil.

pip install psutil
TensorFlow Keras hardware consumption charts

Log hyperparameters

You can log training and model hyperparameters. To do that just pass the parameter dictionary to create_experiment() method:

PARAMS = {'lr': 0.005,
          'momentum': 0.9,
          'epochs': 10,
          'batch_size': 64}

# log params
neptune.create_experiment('keras-tensorflow-more-options', params=PARAMS)
TensorFlow Keras hyperparameter logging

Log model weights

You can log model weights to Neptune both during and after training.

To do that just use a log_artifact() method on the saved model file.

model.save('my_model')

# log model
neptune.log_artifact('my_model')
TensorFlow Keras checkpoints logging

Log image predictions

You can log tensors as images to Neptune with some additional descriptions via log_image() method.

x_test_sample = x_test[:100]
y_test_sample_pred = model.predict(x_test_sample)

for image, y_pred in zip(x_test_sample, y_test_sample_pred):
    description = '\n'.join(['class {}: {}'.format(i, pred)
                                for i, pred in enumerate(y_pred)])
    neptune.log_image('predictions',
                      image,
                      description=description)
TensorFlow Keras logging images

Note

You can log many other experiment metadata like interactive charts, video, audio and more. See the full list.

Remember that you can try it out with zero setup:

How to ask for help?

Please visit the Getting help page. Everything regarding support is there.