keras.MeanSquaredError与自定义reduce_sum(square(diff))损失差异问题

问题现象 #

花费数小时排查以下问题：直接使用keras.MSE可让程序正常运行，但自定义的MSE实现效果偏差很大。原本认为差值平方的均值计算结果应该和keras.MSE非常接近，实际运行时两者初始值差距不大，但后续偏差越来越大，运行日志如下：

step=0 theirs= 13.1761 mine= 14.0251
step=5 theirs= 10.3337 mine= 11.8363
…
step=90 theirs=  0.0361 mine=  6.9888
step=95 theirs=  0.0332 mine=  6.9604

初步排查 #

翻查Keras和TensorFlow源码，发现keras/losses.py中MSE的核心实现是backend.mean(tf.math.squared_difference(y_pred, y_true), axis=-1)，和自定义的tf.reduce_mean(tf.square(y_true-y_pred))逻辑几乎一致，在ipython中测试两者结果也相近，但训练时差异明显，测试代码如下：

import tensorflow as tf
import numpy as np

def small_ds():
    in_t = tf.cast(np.random.randint(5, size=(24, 2)), tf.float32)
    out_t = tf.reduce_sum(in_t, axis=-1)
    return in_t, out_t

def small_model():
    i = tf.keras.layers.Input(shape=(2,))
    d = i
    d = tf.keras.layers.Dense(32, activation="LeakyReLU")(d)
    d = tf.keras.layers.Dense(32, activation="LeakyReLU")(d)
    d = tf.keras.layers.Dense(32, activation="LeakyReLU")(d)
    o = tf.keras.layers.Dense(1, activation="LeakyReLU")(d)
    m = tf.keras.Model(inputs=i, outputs=o)
    return m

def what_is_happening_here():
    opt = tf.keras.optimizers.Adam()
    tf_mse = tf.keras.losses.MeanSquaredError()

    @tf.function
    def my_mse(y_true, y_pred):
        return tf.reduce_mean(tf.square(y_true-y_pred))

    m = small_model()

    @tf.function
    def train_step(x_input, y_true):
        with tf.GradientTape() as tape:
            y_pred = m(x_input, training=True)
            theirs = tf_mse(y_true, y_pred)
            mine   = my_mse(y_true, y_pred)
        grad = tape.gradient(theirs, m.trainable_variables)
        opt.apply_gradients(zip(grad, m.trainable_variables))
        return theirs, mine

    x_input, y_true = small_ds()
    for step in range(100):
        theirs, mine = train_step(x_input, y_true)
        if (step % 5) == 0:
            print(f'step={step} theirs={theirs:8.4f} mine={mine:8.4f}')

if __name__ == '__main__':
    what_is_happening_here()

补充测试 #

最初被第一个回答说服，但后续测试发现完全随机向量测试下，有无归约操作时两者结果完全一致，但放到训练循环中还是会出现异常的非线性差异。怀疑是优化器或计算图有未注意到的逻辑，由于后续需要自定义损失函数，所以需要搞清楚这个问题，也试过把自定义损失继承tf.keras.losses.Loss类，结果还是异常，测试代码如下：

In [22]: tf_mse = tf.keras.losses.MeanSquaredError(reduction=tf.keras.losses.Reduction.NONE)
    ...: my_mse = lambda x,y: tf.reduce_mean(tf.square(x-y), axis=-1)
    ...:
    ...: tf_mser = tf.keras.losses.MeanSquaredError()
    ...: my_mser = lambda x,y: tf.reduce_mean(my_mse(x,y))
    ...:
    ...: y_true = tf.cast(np.random.randint(10, size=(6,1)), tf.float32)
    ...: y_pred = tf.cast(np.random.randint(10, size=(6,1)), tf.float32)
    ...:
    ...: i = tf.keras.layers.Input(shape=(1,))
    ...: o = tf.keras.layers.Dense(32)(i)
    ...: m = tf.keras.Model(inputs=i, outputs=o)
    ...:
    ...: m_pred = m(y_pred)
    ...:
    ...: for a,b in [(tf_mse, tf_mser), (my_mse, my_mser)]:
    ...:     print(f'{a(y_true, y_pred).numpy()} -> {b(y_true, y_pred)}')
    ...:
    ...: for a,b in [(tf_mse, tf_mser), (my_mse, my_mser)]:
    ...:     print(f'{a(y_true, m_pred).numpy()} -> {b(y_true, m_pred)}')
[ 9.  4.  4.  9. 16.  4.] -> 7.666666507720947
[ 9.  4.  4.  9. 16.  4.] -> 7.666666507720947
[20.47 30.15 19.62  9.   12.8  19.62] -> 18.608726501464844
[20.47 30.15 19.62  9.   12.8  19.62] -> 18.608726501464844

根因与解决方案 #

最终找到问题原因，修改small_ds的实现后结果完全一致，修改后的代码如下：

def small_ds():
    in_t = tf.cast(np.random.randint(5, size=(24, 2)), tf.float32)
    out_t = tf.expand_dims(tf.reduce_sum(in_t, axis=-1), -1)
    return in_t, out_t

问题本质是原生MSE会先对最后一个通道(axis=-1)做均值，如果y_true的shape是(24,)，不管归约设置是什么都会按此逻辑计算，给out_t添加expand_dims调整维度后，自定义MSE就和原生实现效果完全一致了。

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