How to Implement Custom Loss Functions and Custom Metrics in TensorFlow?

In deep learning practice, TensorFlow 2.x provides a powerful toolchain for model training and evaluation. However, when default loss functions (e.g., Mean Squared Error (MSE)) or evaluation metrics (e.g., accuracy) fail to meet specific task requirements (e.g., handling imbalanced data, custom business logic, or complex loss structures), custom loss functions and custom metrics become key solutions. This article systematically explains how to implement these features in TensorFlow 2.x, combined with code examples, technical principles, and practical recommendations, ensuring developers can efficiently apply these techniques to enhance model performance.

1. Core Principles of Custom Loss Functions

1. Why Custom Loss Functions Are Needed

Custom loss functions allow developers to:

• Handle non-convex optimization problems
• Integrate business rules
• Implement composite loss

Technical Principles: Loss functions must be differentiable to be compatible with TensorFlow's automatic differentiation mechanism. If the function is not differentiable, training will fail.

TensorFlow recommends implementing by inheriting from the tf.keras.losses.Loss class.

Implementation:

python
import tensorflow as tf

class WeightedMSE(tf.keras.losses.Loss):
    def __init__(self, weights=1.0, name='weighted_mse'):
        super().__init__(name=name)
        self.weights = weights

    def call(self, y_true, y_pred):
        # Calculate squared error and multiply by weights
        error = tf.square(y_true - y_pred)
        return tf.reduce_mean(self.weights * error)

# Usage example: Specify during model compilation
model.compile(optimizer='adam', loss=WeightedMSE(weights=2.0))

Key Notes:

• The weight parameter can be dynamically adjusted (e.g., based on sample importance)
• For sample-level weights (e.g., handling imbalanced data), broadcast the weight tensor to the loss computation.
• Performance optimization: Decorate call with @tf.function to improve execution efficiency:

python
@tf.function
def call(self, y_true, y_pred):
    # Calculate squared error and multiply by weights
    error = tf.square(y_true - y_pred)
    return tf.reduce_mean(self.weights * error)

2. Why Custom Metrics Are Needed

In fraud detection, define F1-score to balance precision and recall. In recommendation systems, compute Recall@K to evaluate recommendation quality. In multi-label classification, implement Jaccard Index.

Technical Principles: Metrics and loss functions are functionally separated: loss is used for optimization, while metrics are used for evaluation; metrics should be gradient-free (i.e., not involved in backpropagation) to avoid training instability.

Custom metrics must inherit from tf.keras.metrics.Metric and implement the following methods:

• reset_states()
• update_states()
• result()

Implementation:

python
import tensorflow as tf

class F1Score(tf.keras.metrics.Metric):
    def __init__(self, name='f1_score', **kwargs):
        super().__init__(name=name, **kwargs)
        self.true_positives = tf.Variable(0.0)
        self.false_positives = tf.Variable(0.0)
        self.false_negatives = tf.Variable(0.0)

    def update_state(self, y_true, y_pred):
        # Update states based on predictions
        y_true = tf.cast(y_true, tf.float32)
        y_pred = tf.cast(y_pred, tf.float32)
        # ... implementation details

    def result(self):
        # Compute F1 score
        precision = self.true_positives / (self.true_positives + self.false_positives + tf.keras.backend.epsilon())
        recall = self.true_positives / (self.true_positives + self.false_negatives + tf.keras.backend.epsilon())
        return 2 * (precision * recall) / (precision + recall + tf.keras.backend.epsilon())

Key Notes:

• Avoid division by zero: Use tf.keras.backend.epsilon() as a safe denominator.
• Handle multi-class: Convert to binary classification using tf.argmax and tf.cast.
• Performance optimization: Use tf.reduce_sum in update_state to avoid loops.

2. Practical Recommendations

1. Loss Function Design Principles

• Ensure the output is a scalar (e.g., tf.reduce_mean), not a tensor.
• Use tf.keras.backend functions (e.g., tf.keras.backend.mean) for framework compatibility.
• Memory management: Avoid creating large temporary tensors in call and use tf.identity.

2. Common Pitfalls

• Loss function is not differentiable: Check if call contains non-differentiable operations (e.g., tf.math.floor), and replace with differentiable functions like tf.math.round.

3. Advanced Techniques

• Directly calling metrics in the loss leads to circular dependencies.
• In call, only compute the loss (e.g., MSE), and register metrics via model.metrics.

4. Final Recommendations

• Prioritize using framework-built-in classes to ensure compatibility.
• Test differentiability: Use tf.test.compute_gradient to verify loss functions.
• Small-batch testing: Validate logic before training using tf.data.Dataset.
• In practical projects, start with simple implementations (e.g., weighted MSE) and gradually extend to complex scenarios (e.g., F1-score).

TensorFlow's documentation and GitHub issues provide rich examples. Recommend combining with source code reading for deeper understanding. Mastering these techniques will significantly enhance model robustness and performance in real-world scenarios.

​