Custom metrics in TensorFlow - ML Experiment: Train & Evaluate

Experiment - Custom metrics

Problem:You have trained a neural network to classify images into two categories. The model shows good accuracy, but you want to track a custom metric called 'balanced accuracy' that accounts for class imbalance.

Current Metrics:Training accuracy: 92%, Validation accuracy: 88%

Issue:The built-in accuracy metric does not reflect how well the model performs on each class equally, especially since the dataset is imbalanced.

Your Task

Add a custom metric to compute balanced accuracy during training and validation, and verify it works correctly.

Use TensorFlow and Keras API only.

Do not change the model architecture or dataset.

Implement the custom metric as a subclass of tf.keras.metrics.Metric.

Hint 1

Hint 2

Hint 3

Solution

TensorFlow

import tensorflow as tf
from tensorflow.keras import layers, models

# Sample dataset: Using MNIST but only digits 0 and 1 for binary classification
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()

# Filter digits 0 and 1
train_filter = (y_train == 0) | (y_train == 1)
test_filter = (y_test == 0) | (y_test == 1)

x_train, y_train = x_train[train_filter], y_train[train_filter]
x_test, y_test = x_test[test_filter], y_test[test_filter]

# Normalize and reshape
x_train = x_train.astype('float32') / 255.0
x_test = x_test.astype('float32') / 255.0
x_train = x_train[..., tf.newaxis]
x_test = x_test[..., tf.newaxis]

# Custom balanced accuracy metric
class BalancedAccuracy(tf.keras.metrics.Metric):
    def __init__(self, name='balanced_accuracy', **kwargs):
        super().__init__(name=name, **kwargs)
        self.recall_0 = tf.keras.metrics.Recall(class_id=0)
        self.recall_1 = tf.keras.metrics.Recall(class_id=1)

    def update_state(self, y_true, y_pred, sample_weight=None):
        # y_pred are logits or probabilities, convert to predicted class
        y_pred_labels = tf.argmax(y_pred, axis=1, output_type=tf.int32)
        self.recall_0.update_state(y_true, y_pred_labels, sample_weight)
        self.recall_1.update_state(y_true, y_pred_labels, sample_weight)

    def result(self):
        recall_0 = self.recall_0.result()
        recall_1 = self.recall_1.result()
        return (recall_0 + recall_1) / 2

    def reset_states(self):
        self.recall_0.reset_states()
        self.recall_1.reset_states()

# Build a simple CNN model
model = models.Sequential([
    layers.Conv2D(16, 3, activation='relu', input_shape=(28, 28, 1)),
    layers.MaxPooling2D(),
    layers.Flatten(),
    layers.Dense(32, activation='relu'),
    layers.Dense(2)  # logits for 2 classes
])

# Compile with custom metric
model.compile(
    optimizer='adam',
    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=['accuracy', BalancedAccuracy()]
)

# Train model
history = model.fit(x_train, y_train, epochs=5, batch_size=64, validation_data=(x_test, y_test))

Created a custom metric class BalancedAccuracy inheriting from tf.keras.metrics.Metric.

Inside the custom metric, used two Recall metrics for each class (0 and 1).

Computed balanced accuracy as the average of recalls for both classes.

Added the custom metric to model.compile metrics list.

Kept model architecture and dataset unchanged.

Added output_type=tf.int32 to tf.argmax to ensure correct dtype for update_state.

Results Interpretation

Before adding the custom metric, only accuracy was tracked, which can hide poor performance on minority classes.

After adding the custom balanced accuracy metric, you can see how well the model performs on each class equally, helping detect bias.

Custom metrics let you track exactly what matters for your problem. Balanced accuracy helps when classes are imbalanced by averaging recall per class.

Bonus Experiment

Try implementing a custom F1-score metric for the same model and dataset.

💡 Hint

Use tf.keras.metrics.Precision and tf.keras.metrics.Recall inside your custom metric to compute F1-score.

Practice

1. What is the main purpose of a custom metric in TensorFlow model training?

easy

A. To track a special score based on true labels and predictions during training

B. To change the model's architecture automatically

C. To speed up the training process by using GPU

D. To save the model weights after each epoch

Solution

1. Step 1: Understand what metrics do

   Metrics measure how well the model performs by comparing true labels and predictions.

2. Step 2: Identify the role of custom metrics

   Custom metrics let you define your own way to measure performance during training.

3. Final Answer:

   To track a special score based on true labels and predictions during training -> Option A

4. Quick Check:

   Custom metric = special score tracking [OK]

Hint: Custom metrics measure model performance uniquely during training [OK]

Common Mistakes:

• Confusing metrics with model architecture changes
• Thinking metrics speed up training
• Believing metrics save model weights

2. Which of the following is the correct way to define a simple custom metric function in TensorFlow?

easy

A. def my_metric(y_true, y_pred): print(y_true, y_pred)

B. def my_metric(y_true, y_pred): return np.mean(np.abs(y_true - y_pred))

C. def my_metric(y_true, y_pred): return y_true + y_pred

D. def my_metric(y_true, y_pred): return tf.reduce_mean(tf.abs(y_true - y_pred))

Solution

1. Step 1: Check TensorFlow compatibility

   Custom metrics must use TensorFlow operations like tf.reduce_mean for compatibility.

2. Step 2: Identify correct function

   def my_metric(y_true, y_pred): return tf.reduce_mean(tf.abs(y_true - y_pred)) uses tf.reduce_mean and tf.abs correctly; def my_metric(y_true, y_pred): return np.mean(np.abs(y_true - y_pred)) uses numpy which is incompatible.

3. Final Answer:

   def my_metric(y_true, y_pred): return tf.reduce_mean(tf.abs(y_true - y_pred)) -> Option D

4. Quick Check:

   Use TensorFlow ops in custom metrics [OK]

Hint: Use TensorFlow ops, not numpy, inside custom metric functions [OK]

Common Mistakes:

• Using numpy functions inside metric causing errors
• Returning tensors without reduction to scalar
• Printing values instead of returning metric

3. Given the following custom metric code, what will be the output metric value for y_true = [1, 2, 3] and y_pred = [1, 2, 4]?

def custom_metric(y_true, y_pred):
    diff = tf.abs(y_true - y_pred)
    return tf.reduce_mean(diff)

medium

A. 1.0

B. 0.6667

C. 0.3333

D. 2.0

Solution

1. Step 1: Calculate absolute differences

   Differences: |1-1|=0, |2-2|=0, |3-4|=1

2. Step 2: Compute mean of differences

   Mean = (0 + 0 + 1) / 3 = 0.3333

3. Final Answer:

   0.3333 -> Option C

4. Quick Check:

   Mean absolute difference = 0.3333 [OK]

Hint: Calculate absolute differences, then average them [OK]

Common Mistakes:

• Summing differences without averaging
• Using squared difference instead of absolute
• Mixing up y_true and y_pred order

4. Identify the error in this custom metric function and choose the correct fix:

def faulty_metric(y_true, y_pred):
    return tf.reduce_sum(y_true - y_pred)

medium

A. Convert y_true and y_pred to numpy arrays first

B. Use tf.reduce_mean instead of tf.reduce_sum to get average error

C. Replace subtraction with addition

D. Add print statements to debug values

Solution

1. Step 1: Understand metric output

   Metrics should return a scalar representing average error, not total sum.

2. Step 2: Fix by averaging

   Replacing tf.reduce_sum with tf.reduce_mean returns average difference, a better metric.

3. Final Answer:

   Use tf.reduce_mean instead of tf.reduce_sum to get average error -> Option B

4. Quick Check:

   Average error metric needs reduce_mean [OK]

Hint: Use reduce_mean for average, not reduce_sum [OK]

Common Mistakes:

• Using reduce_sum causing large values
• Mixing TensorFlow and numpy operations
• Not returning a scalar metric

5. You want to create a custom metric that counts how many predictions differ from true labels by more than 0.5. Which of the following implementations is correct?

hard

A. def count_large_diff(y_true, y_pred):
diff = tf.abs(y_true - y_pred)
return tf.reduce_sum(tf.cast(diff > 0.5, tf.float32))

B. def count_large_diff(y_true, y_pred):
diff = y_true - y_pred
return tf.reduce_mean(diff > 0.5)

C. def count_large_diff(y_true, y_pred):
diff = tf.abs(y_true - y_pred)
return tf.reduce_sum(diff > 0.5)

D. def count_large_diff(y_true, y_pred):
diff = tf.abs(y_true - y_pred)
return tf.reduce_sum(tf.int32(diff > 0.5))

Solution

1. Step 1: Identify correct comparison and casting

   We must compare differences with 0.5 and cast boolean to float for summation.

2. Step 2: Check TensorFlow functions used

   def count_large_diff(y_true, y_pred): diff = tf.abs(y_true - y_pred) return tf.reduce_sum(tf.cast(diff > 0.5, tf.float32)) uses tf.abs, comparison, tf.cast to float32, and tf.reduce_sum correctly.

3. Final Answer:

   def count_large_diff(y_true, y_pred): diff = tf.abs(y_true - y_pred) return tf.reduce_sum(tf.cast(diff > 0.5, tf.float32)) -> Option A

4. Quick Check:

   Cast booleans to float before summing [OK]

Hint: Cast boolean mask to float before summing [OK]

Common Mistakes:

• Not casting boolean mask before summing
• Using reduce_mean instead of reduce_sum
• Using wrong data type cast like tf.int32