TensorFlow模型训练时损失函数出现NaN问题求助

我为人工智能项目开发了一个TensorFlow模型，但训练过程中损失函数出现了NaN值问题。已在训练前检查数据中不存在NaN或无穷值，使用环境为TensorFlow 2.13.0、Python 3.8.10及GPU环境，代码片段如下：

import os
os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0'
import data as dt
print("[Info] Data loaded")

import numpy as np
print("[Info] Numpy loaded")
import random
print("[Info] Random loaded")
import tensorflow as tf
from tensorflow.keras.callbacks import EarlyStopping, LambdaCallback
print("[Info] Tensorflow loaded")
import atexit
print("[Info] Atexit loaded")
import threading
print("[Info] Threading loaded")


# Add this line after importing TensorFlow
physical_devices = tf.config.list_physical_devices('GPU')
if len(physical_devices) > 0:
    tf.config.experimental.set_memory_growth(physical_devices[0], enable=True)

# Limit CPU threads
tf.config.threading.set_intra_op_parallelism_threads(0)
tf.config.threading.set_inter_op_parallelism_threads(0)

inputs = dt.train_inputs
outputs = dt.train_outputs
val_inputs = dt.val_inputs
val_outputs = dt.val_outputs
test_inputs = dt.test_inputs
test_outputs = dt.test_outputs

# Check data for NaNs or infinite values
def check_data(data, name):
    if np.isnan(data).any() or np.isinf(data).any():
        print(f"[Error] {name} contient des NaNs ou des valeurs infinies")
    else:
        print(f"[Info] {name} est valide")

check_data(inputs, "train_inputs")
check_data(outputs, "train_outputs")
check_data(val_inputs, "val_inputs")
check_data(val_outputs, "val_outputs")
check_data(test_inputs, "test_inputs")
check_data(test_outputs, "test_outputs")

class EThAI:
    def __init__(self, save: bool=True, load: bool=dt.values["DoLoad"]):
        # Create a lock to synchronise access to the file
        self.file_access_lock = threading.Lock()
        self.save = bool(save)
        self.load = bool(load)

        self.build_model()

        if self.save:
            self.model.save("CryptoAIModel", save_format='tf')

    def build_model(self, learning_rate=0.0001):
        params = {
            'dense_units1': 256.0,
            'dense_units2': 64
        }
        # Model with optimisated hyperparameter
        self.model = tf.keras.Sequential([
            tf.keras.layers.Input(shape=(33 * 80,)),
            tf.keras.layers.Dense(params["dense_units1"], activation='sigmoid'),
            tf.keras.layers.Dense(params["dense_units2"], activation='sigmoid'),
            tf.keras.layers.Dense(3, activation='softmax')
        ])
        
        if self.load:
            self.model = tf.keras.models.load_model("CryptoAIModel")

        # Compilation of the model with the Adam optimizer and the learning rate planner
        self.model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=learning_rate),
            loss='categorical_crossentropy',
            metrics=['accuracy'])


        print(f"[Info] Model Build (learning_rate={learning_rate})")

    def train(self, epochs: int):
        print("[Info] Starting...")

        def callback(epoch, _):
            self.model.save("CryptoAIModel", save_format='tf')

        early_stopping = EarlyStopping(monitor='val_loss', patience=5, restore_best_weights=True)
        lambda_callback = LambdaCallback(
            on_epoch_end=callback
        )

        tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir="./logs")

        if dt.values["DoBatch"]:
            # Use this list of callbacks when training the model
            self.model.fit(inputs, outputs, epochs=epochs, batch_size=dt.values["BatchSize"], verbose=1, 
                validation_data=(val_inputs, val_outputs),
                callbacks=[early_stopping, lambda_callback, tensorboard_callback])
        else:
            self.model.fit(inputs, outputs, epochs=epochs, batch_size=dt.values["BatchSize"], validation_data=(val_inputs, val_outputs), verbose=1,
                callbacks=[early_stopping, lambda_callback, tensorboard_callback])

        self.model.save("CryptoAIModel", save_format='tf')
        print("[Info] Finished !!!")


    def predict(self, input):
        prediction = self.model.predict(input)
        return prediction
    
    def final_output(self, output):
            values = output[0]
            rounded_values = []
            for value in values:
                if value < 0.5:
                    rounded_values.append(0)
                else:
                    rounded_values.append(1)
            return rounded_values
    
    def get_random_input_output(self, batch_size=dt.values["DoBatch"]):
        num = random.randint(0, len(inputs) - batch_size)
        input_data = np.array([inputs[num+i] for i in range(batch_size)])
        output_data = np.array([outputs[num+i] for i in range(batch_size)])
        return input_data, output_data


if __name__ == '__main__':
    AI = EThAI(save=True)
    try:
        AI.build_model(0.001)
        dt.values["DoLoad"] = True
        AI.load = True
        print("[Info] Training Soon!!!")
        AI.train(20)
        AI.build_model(0.0001)
        AI.train(20)
        AI.build_model(0.00001)
        AI.train(20)
        AI.build_model(0.000001)
        AI.train(20)
        AI.build_model(0.0000001)
        AI.train(20)
        dt.values["DoLoad"] = True
    except Exception as e:
        print(e)
        if AI.save:
            AI.model.save("CryptoAIModel", save_format='tf')

针对训练中损失出现NaN的问题，可从以下几个方向排查和修复：

1. 修正学习率调整逻辑 #

你的代码中连续调用build_model并加载旧模型，会覆盖新设置的学习率，导致学习率调整不生效。正确做法是要么不加载旧模型，要么使用学习率调度器自动调整：

方案A：修改训练逻辑，不重复加载旧模型 #

if __name__ == '__main__':
    AI = EThAI(save=True, load=False)
    try:
        AI.build_model(0.001)
        AI.train(20)
        AI.build_model(0.0001)
        AI.train(20)
        AI.build_model(0.00001)
        AI.train(20)
        AI.build_model(0.000001)
        AI.train(20)
        AI.build_model(0.0000001)
        AI.train(20)
    except Exception as e:
        print(e)
        if AI.save:
            AI.model.save("CryptoAIModel", save_format='tf')

方案B：使用自动学习率调度器 #

from tensorflow.keras.callbacks import ReduceLROnPlateau

# 在train方法中添加调度器
lr_scheduler = ReduceLROnPlateau(monitor='val_loss', factor=0.1, patience=3, min_lr=1e-7)
# 将lr_scheduler加入callbacks列表

2. 标准化输入数据 #

输入数据数值范围过大易引发梯度爆炸，导致NaN。对数据做标准化（均值为0，方差为1）：

from sklearn.preprocessing import StandardScaler

scaler = StandardScaler()
inputs = scaler.fit_transform(inputs)
val_inputs = scaler.transform(val_inputs)
test_inputs = scaler.transform(test_inputs)

3. 优化损失函数的数值稳定性 #

categorical_crossentropy配合softmax输出易出现数值溢出，改用logits输出模式：

# 修改模型输出层，移除softmax激活
tf.keras.layers.Dense(3)

# 编译时设置损失函数为带from_logits=True的版本
self.model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=learning_rate),
                   loss=tf.keras.losses.CategoricalCrossentropy(from_logits=True),
                   metrics=['accuracy'])

4. 添加梯度裁剪 #

限制梯度最大范数，避免梯度爆炸：

# 编译模型时设置梯度裁剪
optimizer = tf.keras.optimizers.Adam(learning_rate=learning_rate, clipnorm=1.0)
self.model.compile(optimizer=optimizer,
                   loss='categorical_crossentropy',
                   metrics=['accuracy'])

5. 修复批处理参数错误 #

get_random_input_output方法的默认batch_size参数为布尔值，会导致数据选取异常，改为使用整数类型的批次大小：

def get_random_input_output(self, batch_size=dt.values["BatchSize"]):
    num = random.randint(0, len(inputs) - batch_size)
    input_data = np.array([inputs[num+i] for i in range(batch_size)])
    output_data = np.array([outputs[num+i] for i in range(batch_size)])
    return input_data, output_data

6. 验证标签格式 #

确保outputs是正确的one-hot编码，每个样本的标签总和为1，格式错误会导致损失计算异常。

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