自研YOLO模型训练无效果且出现NaN，求损失函数优化方案

我正在自研YOLO模型，但训练完全没有效果，还时常因为损失函数中的tf.exp操作出现NaN值。已经参考过YOLOv1、YOLOv2、YOLOv3的论文，以及GitHub上各类公开实现，仍无法解决问题。以下是我的模型结构代码和损失函数代码：

模型结构 #

def res_block(input, filters, size=2, reduse=False, batchnorm=True) :
    skip = layers.Conv2D(filters, (1,1), padding='same')(input)
    for i in range(size) :
        input = layers.Conv2D(filters, 3, padding='same', activation='leaky_relu')(input)
    input = layers.Add()([input, skip])
    input = layers.BatchNormalization()(input)
    input = layers.LeakyReLU(alpha=0.05)(input)
    return input

# 输入
input_img = layers.Input(shape=(IMAGE_SIZE[0], IMAGE_SIZE[1], 3))

x = input_img
x = res_block(x, 64, 2)
x = res_block(x, 64, 2)
x = layers.MaxPooling2D(2)(x)
x = res_block(x, 128, 2)
x = res_block(x, 128, 2)
x = layers.MaxPooling2D(2)(x)
x = res_block(x, 256, 2)
x = res_block(x, 256, 2)
x = layers.MaxPooling2D(2)(x)
x = res_block(x, 256, 2)
x = res_block(x, 256, 2)
x = layers.MaxPooling2D(2)(x)
x = res_block(x, 256, 2)
x = res_block(x, 256, 2)
x = layers.MaxPooling2D(2)(x)
x = res_block(x, 512, 3)

# 输出
output = layers.Conv2D((5+len(MODEL_CLASSES))*len(MODEL_ARCHOR_BOXES), (1, 1), padding='same')(x)
output = layers.Reshape((MODEL_CELLULES[0], MODEL_CELLULES[1], len(MODEL_ARCHOR_BOXES), 5+len(MODEL_CLASSES)))(output)

损失函数 #

"""
步骤1：拆分预测值
"""
# 预测值
pred_x, pred_y , pred_w, pred_h , pred_conf, pred_classes = tf.split(y_pred, (1, 1, 1, 1, 1, len(MODEL_CLASSES)), axis=-1)

# 真实值
true_x , true_y , true_w , true_h , true_conf , true_classes = tf.split(y_true, (1, 1, 1, 1, 1, len(MODEL_CLASSES)), axis=-1)

"""
步骤2：转换预测坐标
"""
converted_pred_x = tf.sigmoid(pred_x)
converted_pred_y = tf.sigmoid(pred_y)
converted_pred_w = tf.exp(pred_w)
converted_pred_h = tf.exp(pred_h)
converted_pred_conf = tf.sigmoid(pred_conf)
converted_pred_classes = tf.math.softmax(pred_classes)

"""
步骤3：计算掩码
"""
iou = calculate_iou(true_x, true_y, true_w, true_h, converted_pred_x, converted_pred_y, converted_pred_w, converted_pred_h)
obj_mask_ij = true_conf

"""
步骤4：计算损失
"""
# XY定位损失
errors_xy = tf.square(converted_pred_x - true_x) + tf.square(converted_pred_y- true_y)
errorsloss_xy = errors_xy * obj_mask_ij
loss_xy = tf.reduce_sum(errorsloss_xy, axis = [0,1,2,3,4])

# WH定位损失
errors_wh = tf.square(tf.sqrt(converted_pred_w)- tf.sqrt(true_w)) + tf.square(tf.sqrt(converted_pred_h) - tf.sqrt(true_h))
errorsloss_wh = errors_wh * obj_mask_ij
loss_wh = tf.reduce_sum(errorsloss_wh, axis = [0,1,2,3,4])

# 含目标的置信度损失
error_conf_obj = tf.square(true_conf - iou * converted_pred_conf)
loss_conf_obj = tf.reduce_sum(error_conf_obj, axis = [0,1,2,3])

# 无目标的置信度损失
error_conf_noobj = tf.square(true_conf - converted_pred_conf)
loss_conf_noobj = tf.reduce_sum(error_conf_noobj, axis = [0,1,2,3])

# 计算总损失
return loss_xy + loss_wh + loss_conf_obj + loss_conf_noobj

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