Monitoring and Optimizing Performance of the Img2Img Model: A Step-by-Step Guide

How to monitor and optimize the performance of the Stable Diffusion Img2Img model using Cerebrium and Arize

Monitoring and Optimizing Performance of the Img2Img Model: A Step-by-Step Guide
Photo by Laura Ockel / Unsplash

In previous posts, we learned about the Img2Img model, saw how to deploy it with Cerebrium, and investigated prompting techniques that will help us get the most out of each generation. In this post, we'll explore how we can use Cerebrium to monitor the performance of the model and discuss some ways to optimize it.

Harnessing the Power of External Monitoring Tools

Cerebrium's functionality extends beyond just model deployment and prediction. It also empowers you to log your model's predictions to external machine-learning monitoring tools. This feature enables real-time performance monitoring, permitting you to compare your model's predictions against actual results and baselines live. Consistent monitoring of model performance is a vital aspect of the machine learning lifecycle, ensuring the model's expected performance aligns with its actual performance.

Presently, Cerebrium supports two monitoring tools - Censius and Arize. In this guide, we'll focus on implementing Arize, but you can find detailed documentation on both Arize and Censius on their respective docs pages.

At a high level, incorporating a monitoring logger into a Conduit object is straightforward. You simply need to call the add_logger method on the object. This method requires the following parameters:

  • platform: The external platform for logging, specified using cerebrium.logging_platform.
  • platform_authentication: A dictionary containing the platform's authentication parameters.
  • features: A list of feature names, each represented as a string.
  • targets: A list of targets, each represented as a string.
  • platform_args: A dictionary of the specific platform's required parameters. Refer to the respective platform's documentation for these details.
  • log_ms: A Boolean value indicating whether the system logs the timestamp in seconds or milliseconds. If True, the system logs the timestamp in milliseconds.

Let's see how this works in practice in the detailed step-by-step walkthrough below.

Adding Monitoring Loggers with Arize

Arize is an ML observability platform that lets you troubleshoot, monitor, and fine-tune your models. In order to monitor the performance of your models using Arize, we first need to add monitoring loggers to the Conduit object in Cerebrium.

First, an account and space on Arize need to be created. The account is basically your profile, while the space is a designated area where you host your specific project.

Next, we introduce the Arize logger to the Conduit object. The Conduit object in Cerebrium is a critical part of the pipeline as it manages the flow of data between your model and various services like Arize.

Let's follow along with the example code in the Cerebrium docs:

from cerebrium import Conduit, model_type, logging_platform
from arize.utils.types import ModelTypes, Environments, Schema, Metrics
features=["pregnancies", "glucose", "bloodPressure", "skinThickness","insulin", "bmi", "diabetesPedigree", "age"]
schema = Schema(

Here we're importing necessary libraries and creating a list of features that our model will be using. Schema is used to define the structure of the data we're sending to Arize. This includes the names of the columns for the prediction id, prediction label, feature, and tags.

platform_args = {
"model_type": ModelTypes.BINARY_CLASSIFICATION,
"schema": schema,
"space_key": "<YOUR_SPACE_KEY>",
"api_key": "<YOUR_API_KEY>",
"features": features,
"metrics_validation": [Metrics.CLASSIFICATION]

In the platform_args dictionary, we specify details about our model and Arize platform like the type of the model, the schema we've defined earlier, space key, API key, features used in the model, and validation metrics.

platform_authentication={"space_key": platform_args['space_key'], "api_key": platform_args['api_key']},

Finally, we add the logger to the Conduit using conduit.add_logger(). We specify the logging platform as Arize and pass in the authentication details. We also specify the features and targets used by the model. platform_args is passed as it is. log_ms=True is used to enable logging of millisecond-level performance.

With all this done, we're ready to deploy our Conduit using conduit.deploy(). After this step, the model starts sending performance data to Arize for real-time monitoring.

Deploying the Conduit with Monitoring Loggers

Once the logger is added to the Conduit object, we need to deploy it. Deploying the Conduit is as simple as calling the conduit.deploy() function. This function initiates the Conduit, thereby activating the flow of data between your model and Arize. With this, Arize starts receiving real-time updates about the model's performance, which can be viewed and analyzed on the Arize platform.

Important Considerations

There are a few key things to keep in mind when working with Arize.

Data Processing Time

When logs are sent to Arize, there is a brief processing period before they become available for analysis. This processing time allows Arize to efficiently handle the data and ensure its accuracy and accessibility.

Logging Actual Values

For each request made to your Cerebrium endpoint, a unique identifier called the run_id is generated. This run_id serves as a crucial link between your model's output and the actual values. By logging the actual values using the run_id, you can effectively compare the performance of your model against baselines and evaluate its accuracy.

Multi-class Classification

In the case of multi-class classification, maintaining consistent class label order is essential. The model returns results based on the defined order of class labels. Let's consider an example using the Iris dataset:

platform_args = {
"model_type": ModelTypes.SCORE_CATEGORICAL,
"schema": schema,
"space_key": "sJlAcWf3",
"api_key": "IB31mRV6AMeiZkuV14zF",
"features": features,
"class_labels": ["Setosa", "Versicolour", "Virginica"]

Suppose we input [5.8, 2.7, 5.1, 1.9] into an XGBoost Classifier. The model's output would be [2], indicating the predicted class label "Virginica" according to the order specified in class_labels.

Handling Embedding Vectors

When dealing with embedding features in your model, it's crucial to handle them separately from the other features. These embedding vectors should not be included in the features array. Instead, define them separately. Here's an example:

features = ['state', 'city', 'merchant_name', 'pos_approved', 'item_count', 'merchant_type', 'charge_amount': 20.11]
embedding_features = {
"nlp_embedding": Embedding(
vector=pd.Series([4.0, 5.0, 6.0, 7.0]),
data="I really like the coffee",
platform_args = {
"model_type": ModelTypes.SCORE_CATEGORICAL,
"schema": schema,
"space_key": "sJlAcWf3",
"api_key": "IB31mRV6AMeiZkuV14zF",
"features": features,
"embedding_features": embedding_features

Remember, when providing input to your model, ensure that the feature values are assigned in the correct order. In the provided example ['ca', 'berkeley', 'Peets Coffee', True, 10, 'coffee shop', 20.11, [4.0, 5.0, 6.0, 7.0]], the values are associated with specific features based on their position in the input list. The embedding feature, represented by [4.0, 5.0, 6.0, 7.0], should be the last entry in the list to maintain proper alignment with the model's expectations.

A Quick Note: Optimizing Model Performance

Let's briefly touch on model optimization. This process begins by identifying any significant divergence between the training and validation loss during the model's training.

If such a divergence is observed, consider applying one or more optimization techniques like modifying the learning rate, adjusting the batch size, or applying different regularization techniques. Check out this guide on optimizing Stable Diffusion for more information.


Monitoring and optimizing the performance of the Img2Img model can significantly enhance its effectiveness. Let's recap the key points:

  1. Importance of Monitoring and Optimization: Monitoring and optimizing the performance of your Img2Img model is crucial for ensuring its effectiveness and identifying areas for improvement.
  2. Cerebrium's Monitoring Capabilities: Cerebrium provides the capability to log your custom model predictions to external ML monitoring tools, allowing real-time performance monitoring.
  3. External Monitoring Tools: We introduced Censius and mentioned the upcoming support for Arize in Q1 2023. These tools provide comprehensive monitoring features to track and analyze your model's performance.
  4. Adding Monitoring Loggers with Cerebrium: By using the add_logger method, you can easily integrate Censius or Arize into your Conduit object, enabling the logging of predictions to the chosen monitoring tool.
  5. Deploying the Conduit with Monitoring Loggers: After adding the appropriate logger to your Conduit object, you can deploy it and start logging predictions to the selected monitoring tool.

Remember, monitoring and optimization are ongoing processes. Stay vigilant, adapt to changes, and unleash the full potential of your Img2Img model with Cerebrium's monitoring tools.

Note: This article originally appeared on the Cerebrium blog.

Subscribe or follow me on Twitter for more content like this!