使用H5格式保存含Balanced Categorical Entropy损失函数的VGG16图像分割模型时加载报错的问题求助

Let me break down why you're hitting this error and how to fix it clearly:

The core issue here is how Keras tracks custom loss functions. Your balanced_cross_entropy is a factory function—it creates and returns the actual loss function (the inner loss function) when you call it with beta. When you compile your model, Keras saves the name of the inner function (loss) to the H5 file, not the name of the factory function.

When you try to load the model with custom_objects={'balanced_cross_entropy(beta)': balanced_cross_entropy(beta)}, Keras is looking for a loss function named balanced_cross_entropy(beta) in the model metadata, but it only finds loss—hence the "Unknown loss function" error.

Solution 1: Assign a Unique Name to Your Loss Function (Recommended) #

Give the inner loss function a distinct, identifiable name that includes your beta parameter. This ensures Keras saves the correct name, and you can map it properly during loading:

beta = 0.5
def balanced_cross_entropy(beta):
    def loss(y_true, y_pred):
        weight_a = beta * tf.cast(y_true, tf.float32)
        weight_b = (1 - beta) * tf.cast(1 - y_true, tf.float32)
        o = (tf.math.log1p(tf.exp(-tf.abs(y_pred))) + tf.nn.relu(-y_pred)) * (weight_a + weight_b) + y_pred * weight_b
        return tf.reduce_mean(o)
    # Set a custom name that includes beta to avoid conflicts
    loss.__name__ = f"balanced_cross_entropy_beta_{beta}"
    return loss

When compiling your model, use the instantiated loss function as before:

model.compile(optimizer='adam', loss=balanced_cross_entropy(beta))

Then, when loading the model, use this custom name to map your loss function:

model = load_model('vgg.h5', custom_objects={f"balanced_cross_entropy_beta_{beta}": balanced_cross_entropy(beta)})

Solution 2: Map the Default "loss" Name (Quick Fix, Less Reliable) #

If you don't want to modify your loss function definition, you can directly map the loss name (which Keras saved) to your instantiated loss function. This works for simple cases but can cause conflicts if you have multiple custom loss functions:

model = load_model('vgg.h5', custom_objects={'loss': balanced_cross_entropy(beta)})

Extra Tip for Robustness #

For more advanced use cases (like loading models with different beta values or sharing models across environments), you can implement a class-based loss function with a get_config method. This lets Keras save the beta parameter along with the loss function, so you don't have to hardcode it during loading:

import tensorflow as tf
from tensorflow.keras.losses import Loss

class BalancedCrossEntropy(Loss):
    def __init__(self, beta=0.5, name="balanced_cross_entropy"):
        super().__init__(name=name)
        self.beta = beta

    def call(self, y_true, y_pred):
        weight_a = self.beta * tf.cast(y_true, tf.float32)
        weight_b = (1 - self.beta) * tf.cast(1 - y_true, tf.float32)
        o = (tf.math.log1p(tf.exp(-tf.abs(y_pred))) + tf.nn.relu(-y_pred)) * (weight_a + weight_b) + y_pred * weight_b
        return tf.reduce_mean(o)

    def get_config(self):
        config = super().get_config()
        config.update({"beta": self.beta})
        return config

Then compile and load like this:

# Compile
model.compile(optimizer='adam', loss=BalancedCrossEntropy(beta=0.5))
# Load
model = load_model('vgg.h5', custom_objects={'BalancedCrossEntropy': BalancedCrossEntropy})

This class-based approach is cleaner for parameterized losses and avoids name mapping issues entirely.

内容的提问来源于stack exchange，提问作者GabS