求Keras代码的等效TensorFlow实现，解决Loss无法降至0.011以下问题

Hey there! I get it—when you have a Keras model that's performing great (loss nearly hitting 0) but your native TensorFlow implementation can't get loss below 0.011, it's usually down to missing subtle details that Keras handles under the hood. Let's walk through creating an equivalent TF implementation that matches your Keras model's behavior step by step.

First, let's start with a sample well-performing Keras model that aligns with your xData and yTrainData setup (adjust layers/hyperparameters to match your actual Keras code):

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

# Sample Keras model (match your actual model structure!)
keras_model = keras.Sequential([
    layers.Dense(64, activation='relu', input_shape=(xData.shape[1],)),
    layers.Dense(32, activation='relu'),
    # Adjust output layer based on task: 1 for regression, num_classes for classification
    layers.Dense(1 if yTrainData.ndim == 1 else yTrainData.shape[1], activation='linear')
])

# Compile with matching optimizer/loss
keras_model.compile(
    optimizer=keras.optimizers.Adam(learning_rate=1e-3),
    loss=keras.losses.MeanSquaredError() if yTrainData.ndim == 1 else keras.losses.CategoricalCrossentropy(from_logits=True)
)

# Train (matches typical Keras fit setup)
history = keras_model.fit(
    xData, yTrainData,
    batch_size=32,
    epochs=100,
    validation_split=0.1
)

This implementation replicates every key behavior of the Keras model to ensure you get the same low-loss performance:

import tensorflow as tf

# Set random seed for reproducibility (matches Keras default)
tf.random.set_seed(42)

# 1. Define model class (matches Keras Sequential structure exactly)
class TFModel(tf.keras.Model):
    def __init__(self, input_dim, output_dim):
        super().__init__()
        # Use Keras's default initializers for Dense layers
        self.dense1 = tf.keras.layers.Dense(
            64, activation='relu',
            kernel_initializer='glorot_uniform', bias_initializer='zeros'
        )
        self.dense2 = tf.keras.layers.Dense(
            32, activation='relu',
            kernel_initializer='glorot_uniform', bias_initializer='zeros'
        )
        self.dense3 = tf.keras.layers.Dense(
            output_dim, activation='linear',
            kernel_initializer='glorot_uniform', bias_initializer='zeros'
        )
    
    def call(self, inputs, training=False):
        x = self.dense1(inputs)
        x = self.dense2(x)
        return self.dense3(x)

# Calculate input/output dimensions from your data
input_dim = xData.shape[1]
output_dim = 1 if yTrainData.ndim == 1 else yTrainData.shape[1]

# Initialize the model
model = TFModel(input_dim, output_dim)

# 2. Match optimizer and loss function exactly
if yTrainData.ndim == 1:
    # Regression task
    loss_fn = tf.keras.losses.MeanSquaredError()
else:
    # Classification with one-hot labels (adjust if using sparse labels!)
    loss_fn = tf.keras.losses.CategoricalCrossentropy(from_logits=True)

optimizer = tf.keras.optimizers.Adam(learning_rate=1e-3)

# 3. Prepare dataset to match Keras fit behavior
batch_size = 32
val_split = 0.1
val_size = int(len(xData) * val_split)

# Split data into train/validation (same as Keras's validation_split)
x_train, x_val = xData[:-val_size], xData[-val_size:]
y_train, y_val = yTrainData[:-val_size], yTrainData[-val_size:]

# Create shuffled dataset (Keras shuffles training data every epoch by default)
train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
train_dataset = train_dataset.shuffle(buffer_size=len(x_train)).batch(batch_size)

val_dataset = tf.data.Dataset.from_tensor_slices((x_val, y_val)).batch(batch_size)

# 4. Training loop with GradientTape (replicates Keras's internal training step)
epochs = 100
train_loss_history = []
val_loss_history = []

for epoch in range(epochs):
    epoch_train_loss = 0.0
    epoch_val_loss = 0.0
    train_batches = 0
    val_batches = 0

    # Training step
    for x_batch, y_batch in train_dataset:
        with tf.GradientTape() as tape:
            # Enable training mode (critical for layers like Dropout/BatchNorm)
            y_pred = model(x_batch, training=True)
            loss = loss_fn(y_batch, y_pred)
        
        # Update weights
        gradients = tape.gradient(loss, model.trainable_variables)
        optimizer.apply_gradients(zip(gradients, model.trainable_variables))
        
        epoch_train_loss += loss.numpy()
        train_batches += 1
    
    # Average training loss for the epoch
    avg_train_loss = epoch_train_loss / train_batches
    train_loss_history.append(avg_train_loss)

    # Validation step
    for x_val_batch, y_val_batch in val_dataset:
        y_val_pred = model(x_val_batch, training=False)
        val_loss = loss_fn(y_val_batch, y_val_pred)
        epoch_val_loss += val_loss.numpy()
        val_batches += 1
    
    avg_val_loss = epoch_val_loss / val_batches
    val_loss_history.append(avg_val_loss)

    # Print progress (matches Keras's fit output)
    print(f"Epoch {epoch+1}/{epochs} | Train Loss: {avg_train_loss:.6f} | Val Loss: {avg_val_loss:.6f}")

Critical Details to Ensure Matching Performance #

• Initializers: Keras uses glorot_uniform (Xavier uniform) for Dense layer kernels and zeros for biases by default—don't skip this, as bad initialization can stall training.
• Training Mode: Always set training=True during training and False during validation (this affects layers like Dropout or BatchNormalization that behave differently in each mode).
• Dataset Shuffling: Keras shuffles the entire training dataset every epoch—our shuffle(buffer_size=len(x_train)) replicates this.
• Loss Function: Double-check you're using the exact same loss as your Keras model (e.g., if you used from_logits=True in Keras, don't omit it in TF).
• Optimizer Hyperparameters: Match learning rate, momentum, and any other optimizer settings exactly to your Keras code.

If your original Keras model includes other layers (like Dropout, BatchNormalization, or Conv2D), just add them to the TFModel class and handle the training parameter in the call method—this implementation is designed to be easily extended to match your specific model.

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