TensorFlow

Integrate W&B with TensorFlow for logging custom metrics, using estimator hooks, and TensorBoard log synchronization.

This page shows how to integrate W&B with TensorFlow to track experiments, log metrics, and synchronize TensorBoard logs. Follow these patterns to capture training data from TensorFlow models, customize what you log through estimator hooks or manual logging, and reuse your existing TensorBoard workflows with W&B’s centralized dashboard. This page targets TensorFlow users who want richer experiment tracking than TensorBoard alone provides.

Get started

If you already use TensorBoard, you can integrate with W&B. Import both libraries to make the W&B and TensorFlow APIs available in your script.

import tensorflow as tf
import wandb

Log custom metrics

This section covers how to log metrics that TensorBoard doesn’t already capture, so you can track additional values alongside your standard TensorBoard summaries.

If you need to log additional custom metrics that TensorBoard doesn’t log, you can call run.log() in your code, for example run.log({"custom": 0.8}).

W&B turns off the step argument in run.log() when syncing TensorBoard. To set a different step count, log the metrics with a step metric as:

with wandb.init(config=tf.flags.FLAGS, sync_tensorboard=True) as run:
    run.log({"custom": 0.8, "global_step":global_step}, step=global_step)

TensorFlow estimators hook

This section describes the W&B hook for TensorFlow estimators, which gives you fine-grained control over what W&B captures during estimator training.

If you want more control over what gets logged, W&B also provides a hook for TensorFlow estimators. It logs all tf.summary values in the graph.

import tensorflow as tf
import wandb

run = wandb.init(config=tf.FLAGS)

estimator.train(hooks=[wandb.tensorflow.WandbHook(steps_per_log=1000)])
run.finish()

Log manually

If you’re not using estimators or want to log specific summaries explicitly, this section shows how to send tf.summary values to W&B directly.

One way to log metrics in TensorFlow is to log tf.summary with the TensorFlow logger:

import wandb
run = wandb.init(config=tf.flags.FLAGS, sync_tensorboard=True)
with tf.Session() as sess:
    # ...
    wandb.tensorflow.log(tf.summary.merge_all())

With TensorFlow 2, the recommended way to train a model with a custom loop is to use tf.GradientTape. For more information, see the TensorFlow custom training walkthrough. To incorporate W&B to log metrics in your custom TensorFlow training loops, follow this snippet:

    with tf.GradientTape() as tape:
        # Get the probabilities
        predictions = model(features)
        # Calculate the loss
        loss = loss_func(labels, predictions)

    # Log your metrics
    run.log("loss": loss.numpy())
    # Get the gradients
    gradients = tape.gradient(loss, model.trainable_variables)
    # Update the weights
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))

A full example of customizing training loops in TensorFlow 2 is available.

Differences between W&B and TensorBoard

If you’re evaluating whether to adopt W&B alongside or in place of TensorBoard, this section highlights the key differences.

W&B was built to address common limitations TensorBoard users encountered. Here are areas where W&B differs:

• Reproduce models: W&B supports experimentation, exploration, and reproducing models later. W&B captures metrics, hyperparameters, and the code version, and can save your version-control status and model checkpoints so your project is reproducible.
• Automatic organization: When you’re picking up a project from a collaborator, returning after time away, or revisiting an old project, W&B lets you see the models you’ve tried so you don’t re-run experiments unnecessarily.
• Flexible integration: Add W&B to your project by installing the open-source Python package and adding a few lines to your code. Each run produces logged metrics and records.
• Persistent, centralized dashboard: Whether you train your models on your local machine, a shared lab cluster, or spot instances in the cloud, W&B sends your results to the same centralized dashboard. You don’t need to copy and organize TensorBoard files from different machines.
• Tables: Search, filter, sort, and group results from different models. You can review model versions and find the best-performing models for different tasks.
• Tools for collaboration: Use W&B to organize machine learning projects. Share a link to W&B, or use private teams to send results to a shared project. Reports support collaboration through interactive visualizations and Markdown descriptions, which you can use to keep a work log, share findings with your supervisor, or present findings to your lab or team.

To try W&B, create a free account.

Examples

To see these integration patterns applied to complete projects, explore the following examples:

• MNIST example using TensorFlow Estimators.
• Fashion MNIST example using raw TensorFlow.
• W&B dashboard for the MNIST example.
• Customizing training loops in TensorFlow 2: article and dashboard.