训练期间的 Tensorflow 自定义损失函数 NaN

Tensorflow custom loss function NaNs during training

我正在尝试为检测对象并在其上绘制非轴对齐边界框的 CNN 编写自定义损失函数。我的输入是 200x200 图像,输出/标签是形式为

的 6d 向量
[object_present, x,y, angle, width, height]

其中object_present是表示对象是否存在的二进制特征,(x,y)是边界框的中心,angle是bbox从轴对齐的旋转角度, width 和 height 是 bbox 的尺寸。当 object_present = 0 时,所有其他特征都设置为 NaN。

因此,我的自定义损失函数需要忽略负样本的 NaN,并将二元交叉熵损失应用于 object_present 特征。对于正样本,我还必须包括 (x,y) 和宽度、高度的 MSE 损失,以及我定义为 arctan(sin(angle1 - angle2), cos(angle1-angle2) 的 angular 回归损失)).我的实现如下:

binary_loss_func = tf.keras.losses.BinaryCrossentropy()
def loss_func(true_labels, pred_labels):
    binary_loss = binary_loss_func(true_labels[:,0], pred_labels[:, 0])
    mse_loss1 = tf.reduce_mean(tf.where(tf.math.is_nan(true_labels[:,1:3]), tf.zeros_like(true_labels[:, 1:3]),
    tf.square(tf.subtract(true_labels[:, 1:3], pred_labels[:, 1:3])))) 
    mse_loss2 = tf.reduce_mean(tf.where(tf.math.is_nan(true_labels[:,4:]), 
    tf.zeros_like(true_labels[:, 4:]), tf.square(tf.subtract(true_labels[:, 4:], pred_labels[:, 4:]))))
    angular_loss = tf.reduce_mean(tf.where(is_nan(true_labels[:,3]), tf.zeros_like(true_labels[:, 3]), 
    tf.abs(tf.atan2(tf.sin(true_labels[:, 3] - pred_labels[:, 3]), tf.cos(true_labels[:, 3] - pred_labels[:, 3])))))
    return  mse_loss1 + mse_loss2 + binary_loss + angular_loss

我的问题是在第一批训练后这个 returns NaN 损失值(只有第一批不给 NaN 损失),尽管我认为上面的代码应该 return负样本0损失。我已经确认 Binary Loss 函数是 returning 实数,所以问题出在损失的其他部分。在使用 tf.print 语句进行一些调试后,我发现 pred_labels 在第一批训练后变为 NaN。我不确定为什么会发生这种情况,以及这是否与我的自定义损失函数的定义方式有关,或者是否与我的模型有关。我使用的模型是:

IMAGE_SIZE = 200
CONV_PARAMS = {"kernel_size": 3, "use_bias": False, "padding": "same"}
CONV_PARAMS2 = {"kernel_size": 5, "use_bias": False, "padding": "same"}

model = Sequential()
model.add(
    Reshape((IMAGE_SIZE, IMAGE_SIZE, 1), input_shape=(IMAGE_SIZE, IMAGE_SIZE))
)
model.add(Conv2D(16, **CONV_PARAMS))
model.add(BatchNormalization())
model.add(Activation('relu'))
model.add(MaxPool2D())
model.add(Conv2D(32, **CONV_PARAMS))
model.add(BatchNormalization())
model.add(Activation('relu'))
model.add(MaxPool2D())
model.add(Conv2D(64, **CONV_PARAMS))
model.add(BatchNormalization())
model.add(Activation('relu'))
model.add(MaxPool2D())
model.add(Flatten())
model.add(Dense(6))

您似乎仍在使用 nan 值计算您的损失,尽管您正在努力避免它。也许尝试这样的事情:

binary_loss_func = tf.keras.losses.BinaryCrossentropy()
def loss_func(true_labels, pred_labels):
    true_labels = tf.where(tf.math.is_nan(true_labels), tf.zeros_like(true_labels), true_labels)
    condition = tf.equal(true_labels, 0.0)

    binary_loss = tf.where(condition, tf.reduce_mean(true_labels), binary_loss_func(true_labels[:,0], pred_labels[:, 0]))

    mse_loss1 = tf.reduce_mean(tf.where(tf.equal(true_labels[:, 1:3], 0.0), true_labels[:, 1:3],
                                        tf.square(tf.subtract(true_labels[:, 1:3], pred_labels[:, 1:3])))) 

    mse_loss2 = tf.reduce_mean(tf.where(tf.equal(true_labels[:, 4:], 0.0), true_labels[:, 4:], 
                                        tf.square(tf.subtract(true_labels[:, 4:], pred_labels[:, 4:]))))

    angular_loss = tf.reduce_mean(tf.where(tf.equal(true_labels[:, 3], 0.0), true_labels[:, 3], 
    tf.abs(tf.atan2(tf.sin(true_labels[:, 3] - pred_labels[:, 3]), tf.cos(true_labels[:, 3] - pred_labels[:, 3])))))

    return  mse_loss1 + mse_loss2 + binary_loss + angular_loss