Skip to content

Integrate with TensorFlow

TensorFlow is an end-to-end open source platform for machine learning. It lets researchers push the state-of-the-art in machine learning and developers easily build and deploy machine learning-powered applications.

Instrument TensorFlow with Comet to start managing experiments, create dataset versions and track hyperparameters for faster and easier reproducibility and collaboration.

Note: If you are using the Keras API with TensorFlow, see our Keras Integration.

Comet also integrates with TensorFlow's TensorBoard.

Start logging

Connect Comet to your existing code by adding in a simple Comet Experiment.

Add the following lines of code to your script or notebook:

import comet_ml
import tensorflow as tf
experiment = comet_ml.Experiment(
    api_key="<Your API Key>",
    project_name="<Your Project Name>"
)
# Your code here

Note

There are other ways to configure Comet. See more here.

Log automatically

After an Experiment has been created, Comet automatically logs the following TensorFlow items, by default, with no additional configuration:

  • Model and graph description
  • Steps and epochs
  • Metrics (such as loss and accuracy)
  • Hyperparameters
  • Optimizer Parameters
  • Number of trainable parameters

Note

Don't see what you need to log here? We have your back. You can manually log any kind of data to Comet using the Experiment object. For example, use experiment.log_image to log images, or experiment.log_audio to log audio.

Configure Comet for TensorFlow

You can further configure what is automatically logged by Comet through an Experiment parameter, an environment variable, or a configuration setting:

ItemExperiment ParameterEnvironment SettingConfiguration Setting
model/graph descriptionlog_graphCOMET_AUTO_LOG_GRAPHcomet.auto_log.graph
metricsauto_metric_loggingCOMET_AUTO_LOG_METRICScomet.auto_log.metrics
metric logging rateauto_metric_step_rateCOMET_AUTO_LOG_METRIC_STEP_RATEcomet.auto_log.metric_step_rate
hyperparametersauto_param_loggingCOMET_AUTO_LOG_PARAMETERScomet.auto_log.parameters
command-line argumentsparse_argsCOMET_AUTO_LOG_CLI_ARGUMENTScomet.auto_log.cli_arguments

Integrate with TensorFlow's TensorBoard

By integrating Comet with TensorBoard, all data that is written to TensorBoard is automatically logged to Comet.

End-to-end example

import comet_ml
import tensorflow as tf

from tensorflow.keras.layers import Dense, Flatten, Conv2D
from tensorflow.keras import Model

experiment = comet_ml.Experiment()
writer = tf.summary.create_file_writer("./mylogs")

class MyModel(Model):
    def __init__(self):
        super(MyModel, self).__init__()
        self.conv1 = Conv2D(32, 3, activation="relu")
        self.flatten = Flatten()
        self.d1 = Dense(128, activation="relu")
        self.d2 = Dense(10)

    def call(self, x):
        x = self.conv1(x)
        x = self.flatten(x)
        x = self.d1(x)
        return self.d2(x)


# Create an instance of the model
model = MyModel()

mnist = tf.keras.datasets.mnist

(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0

# Add a channels dimension
x_train = x_train[..., tf.newaxis].astype("float32")
x_test = x_test[..., tf.newaxis].astype("float32")

train_ds = tf.data.Dataset.from_tensor_slices(
    (x_train, y_train)).shuffle(10000).batch(32)

test_ds = tf.data.Dataset.from_tensor_slices((x_test, y_test)).batch(32)

optimizer = tf.keras.optimizers.Adam()
loss_object = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)

train_loss = tf.keras.metrics.Mean(name='train_loss')
train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='train_accuracy')

test_loss = tf.keras.metrics.Mean(name='test_loss')
test_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='test_accuracy')

@tf.function
def train_step(images, labels):
    with tf.GradientTape() as tape:
        # training=True is only needed if there are layers with different
        # behavior during training versus inference (e.g. Dropout).
        predictions = model(images, training=True)
        loss = loss_object(labels, predictions)
    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))

    train_loss(loss)
    train_accuracy(labels, predictions)

@tf.function
def test_step(images, labels):
    # training=False is only needed if there are layers with different
    # behavior during training versus inference (e.g. Dropout).
    predictions = model(images, training=False)
    t_loss = loss_object(labels, predictions)

    test_loss(t_loss)
    test_accuracy(labels, predictions)

EPOCHS = 2

for epoch in range(EPOCHS):
    # Reset the metrics at the start of the next epoch
    train_loss.reset_states()
    train_accuracy.reset_states()
    test_loss.reset_states()
    test_accuracy.reset_states()

    for images, labels in train_ds:
        train_step(images, labels)
    for test_images, test_labels in test_ds:
        test_step(test_images, test_labels)

    with writer.as_default():
        tf.summary.scalar("train_loss", train_loss.result(), step=epoch)
        tf.summary.scalar("train_accuracy", train_accuracy.result(), step=epoch)
        tf.summary.scalar("test_loss", test_loss.result(), step=epoch)
        tf.summary.scalar("test_accuracy", test_accuracy.result(), step=epoch)

    print(
        f"Epoch {epoch + 1}, "
        f"Loss: {train_loss.result()}, "
        f"Accuracy: {train_accuracy.result() * 100}, "
        f"Test Loss: {test_loss.result()}, "
        f"Test Accuracy: {test_accuracy.result() * 100}"
    )

Try out Comet with TensorBoard!

Try out this example Colab Notebook for using Comet with TensorBoard.

Open In Colab

Aug. 3, 2022