自定义Keras模型使用自定义损失函数训练时报错的解决咨询

Let's break down exactly what's going wrong here and walk through the fixes step by step:

1. Your call Method Is Defined Incorrectly #

Right now, your call method doesn't accept any input parameters, and you're trying to use a pre-defined Input layer from __init__—that's not how Keras subclass models work. The call method should be the entry point for your input tensors; you don't need to define an Input layer in the initialization step.

2. Your Custom compile Method Skips Critical Keras Logic #

You're overriding compile but only calling super().compile() without passing parameters, then manually setting self.optimizer and self.loss. This skips Keras's internal setup for components like compiled_loss and compiled_metrics, which is why the training loop can't find the loss parameter later.

Fixed Full Code #

Here's the corrected version with explanations in comments:

import numpy as np
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import backend as K
from numpy.random import seed

class CustomModel(keras.Model):
    def __init__(self, b, input_dim):
        super(CustomModel, self).__init__()
        # Removed the Input layer definition—we'll handle inputs directly in call()
        self.dense1 = keras.layers.Dense(
            20, 
            name='hidden', 
            kernel_initializer=initializer, 
            bias_initializer=initializer, 
            activation=lambda x: K.tanh(b * x)
        )
        self.dense2 = keras.layers.Dense(
            2, 
            activation='linear', 
            name='output', 
            use_bias=False, 
            trainable=False,
            kernel_initializer=lambda shape, dtype: initializeOutputWeights(shape, dtype)
        )
        # Initialize our accuracy metric directly in the model
        self.acc_metric = keras.metrics.SparseCategoricalAccuracy(name="accuracy")

    # Fixed call method: now accepts input tensors as a parameter
    def call(self, inputs):
        x1 = self.dense1(inputs)
        return self.dense2(x1)

    # Fixed compile method: pass params to parent class so Keras handles internal setup
    def compile(self, optimizer, loss, **kwargs):
        super(CustomModel, self).compile(optimizer=optimizer, loss=loss, **kwargs)
        # No need to manually set self.loss/optimizer—parent compile handles this properly

    def train_step(self, data):
        x, y = data
        with tf.GradientTape() as tape:
            y_pred = self(x, training=True)  # Forward pass
            # Use self.compiled_loss instead of direct self.loss call—this handles regularization losses automatically
            loss = self.compiled_loss(y, y_pred, regularization_losses=self.losses)
        
        # Compute and apply gradients
        trainable_vars = self.trainable_variables
        gradients = tape.gradient(loss, trainable_vars)
        self.optimizer.apply_gradients(zip(gradients, trainable_vars))
        
        # Update our accuracy metric
        self.acc_metric.update_state(y, y_pred)
        # Return metrics for training logs
        return {
            'loss': loss,
            self.acc_metric.name: self.acc_metric.result()
        }

def initializeOutputWeights(shape, dtype=None):
    randoms = np.random.randint(low=2, size=shape)
    new = np.where(randoms == 0, -1, randoms)
    return K.variable(new, dtype=dtype)

class customLoss(keras.losses.Loss):
    def __init__(self, d=10, name="CustomLoss"):
        super().__init__(name=name)
        self.d = d

    def call(self, y_true, y_pred):
        # Use tf.shape instead of len() for graph-mode compatibility
        N = tf.shape(y_true)[0]
        L = tf.shape(y_pred)[1]
        y_dot = y_pred * y_true
        y_d = y_dot - self.d
        y_square = y_d * y_d
        # Simplified masking with tf.where (cleaner and more efficient than tensor_scatter_nd_update)
        y_loss = tf.where(y_dot > self.d, 0.0, y_square)
        return tf.divide(tf.reduce_sum(y_loss), tf.cast(N * L, tf.float32))

# Seed setup for reproducibility
seed(1)
tf.random.set_seed(2)

initializer = tf.keras.initializers.RandomUniform(minval=-1, maxval=1)
b = np.ones(20)
cModel = CustomModel(b, 9)
Losscustom = customLoss(d=16)

# Compile works normally now
cModel.compile(optimizer='adam', loss=Losscustom)

# Fit your model once X_train/y_train are defined
# cModel.fit(X_train, y_train, batch_size=64, epochs=2)

Key Fix Details #

• call Method: Now accepts inputs and processes them directly—this is the standard way to build subclassed Keras models.
• compile Method: We pass all parameters to the parent compile method, letting Keras handle the internal setup for loss, optimizer, and metrics.
• Loss Function: Swapped tensor_scatter_nd_update for tf.where (cleaner and faster) and used tf.shape instead of len() to work properly in TensorFlow's graph mode.
• train_step: Using self.compiled_loss instead of calling self.loss directly ensures regularization losses are included automatically.

For Future Custom Optimization #

Since you mentioned needing custom optimization later, you can still modify the train_step method to replace self.optimizer.apply_gradients with your own gradient update logic—this fix doesn't restrict that at all. Just keep the gradient calculation and weight update logic consistent with TensorFlow's eager/graph mode rules.

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