自定义皮尔森损失函数遇到错误：tensorflow.python.framework.errors_impl.InvalidArgumentError

出错

tensorflow.python.framework.errors_impl.InvalidArgumentError: 2 root error(s) found.
  (0) Invalid argument:  Incompatible shapes: [20,1335] vs. [26700]
	 [[node PCCLoss/sub_1 (defined at D:\ProjectWork\Tensorflow2.0Work

自定义损失函数

class PCCLoss(losses.Loss):
    def __init__(self, n):
        super(PCCLoss,self).__init__()
        self.n = n

    def call(self,y_true,y_pred):
        # n = y_pred.shape[-1]
        sum_xy = tf.math.reduce_sum(y_true * y_pred, axis=1)
        sum_x = tf.math.reduce_sum(tf.reshape(tf.math.reduce_sum(y_true, axis=1), (-1,1)), axis=1)
        sum_y = tf.math.reduce_sum(tf.reshape(tf.math.reduce_sum(y_pred, axis=1), (-1,1)), axis=1)
        sum_x2 = tf.math.reduce_sum(y_true * y_true, axis=1)
        sum_y2 = tf.math.reduce_sum(y_pred * y_pred, axis=1)
        pcc = (self.n * sum_xy - sum_x * sum_y) / tf.math.sqrt((self.n * sum_x2 - sum_x * sum_x) * (self.n * sum_y2 - sum_y * sum_y))

        return tf.math.reduce_mean(pcc)

使用自定义损失

model.compile(optimizer=keras.optimizers.Adam(LR),
             loss=PCCLoss(1335),
             metrics=['accuracy'])

具体错误如下：

Epoch 1/10
2020-10-03 20:40:26.807772: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library cublas64_10.dll
Traceback (most recent call last):
  File "D:/ProjectWork/Tensorflow2.0Work/Spectrum2spectrum922/maintrain01.py", line 149, in <module>
    callbacks=[early_stop, check_callback, lr_schedule])
  File "C:\Users\OFC\Anaconda3\envs\tf2.0\lib\site-packages\tensorflow\python\keras\engine\training.py", line 108, in _method_wrapper
    return method(self, *args, **kwargs)
  File "C:\Users\OFC\Anaconda3\envs\tf2.0\lib\site-packages\tensorflow\python\keras\engine\training.py", line 1098, in fit
    tmp_logs = train_function(iterator)
  File "C:\Users\OFC\Anaconda3\envs\tf2.0\lib\site-packages\tensorflow\python\eager\def_function.py", line 780, in __call__
    result = self._call(*args, **kwds)
  File "C:\Users\OFC\Anaconda3\envs\tf2.0\lib\site-packages\tensorflow\python\eager\def_function.py", line 840, in _call
    return self._stateless_fn(*args, **kwds)
  File "C:\Users\OFC\Anaconda3\envs\tf2.0\lib\site-packages\tensorflow\python\eager\function.py", line 2829, in __call__
    return graph_function._filtered_call(args, kwargs)  # pylint: disable=protected-access
  File "C:\Users\OFC\Anaconda3\envs\tf2.0\lib\site-packages\tensorflow\python\eager\function.py", line 1848, in _filtered_call
    cancellation_manager=cancellation_manager)
  File "C:\Users\OFC\Anaconda3\envs\tf2.0\lib\site-packages\tensorflow\python\eager\function.py", line 1924, in _call_flat
    ctx, args, cancellation_manager=cancellation_manager))
  File "C:\Users\OFC\Anaconda3\envs\tf2.0\lib\site-packages\tensorflow\python\eager\function.py", line 550, in call
    ctx=ctx)
  File "C:\Users\OFC\Anaconda3\envs\tf2.0\lib\site-packages\tensorflow\python\eager\execute.py", line 60, in quick_execute
    inputs, attrs, num_outputs)
tensorflow.python.framework.errors_impl.InvalidArgumentError: 2 root error(s) found.
  (0) Invalid argument:  Incompatible shapes: [20,1335] vs. [26700]
	 [[node PCCLoss/sub_1 (defined at D:\ProjectWork\Tensorflow2.0Work\Spectrum2spectrum922\pcc_loss_function.py:150) ]]
	 [[gradient_tape/Multi_model/embedding_1/embedding_lookup/Reshape_1/_48]]
  (1) Invalid argument:  Incompatible shapes: [20,1335] vs. [26700]
	 [[node PCCLoss/sub_1 (defined at D:\ProjectWork\Tensorflow2.0Work\Spectrum2spectrum922\pcc_loss_function.py:150) ]]
0 successful operations.
0 derived errors ignored. [Op:__inference_train_function_8818]

Function call stack:
train_function -> train_function


Process finished with exit code 1

原因
sum_xy = tf.math.reduce_sum(y_true * y_pred, axis=1)
其中 axis=1，因为我的输入数据是shape=(1000, 1335),二维，按行（row）进行加（sum）运算。但是数据在tensorflow数据流中会改变shape,所以为了能够继续按照行运算，设置axis=-1。
正确运行的代码：

class PCCLoss(losses.Loss):
    def __init__(self, n):
        super(PCCLoss,self).__init__()
        self.n = n

    def call(self,y_true,y_pred):
        # n = y_pred.shape[-1]
        sum_xy = tf.math.reduce_sum(y_true * y_pred, axis=-1)
        sum_x = tf.math.reduce_sum(tf.reshape(tf.math.reduce_sum(y_true, axis=-1), (-1,1)), axis=-1)
        sum_y = tf.math.reduce_sum(tf.reshape(tf.math.reduce_sum(y_pred, axis=-1), (-1,1)), axis=-1)
        sum_x2 = tf.math.reduce_sum(y_true * y_true, axis=-1)
        sum_y2 = tf.math.reduce_sum(y_pred * y_pred, axis=-1)
        conv= self.n * sum_xy - sum_x * sum_y
        var1 = self.n * sum_x2 - sum_x * sum_x
        var2 = self.n * sum_y2 - sum_y * sum_y
        pcc = conv / tf.math.sqrt(var1*var2)

        return 10*(1-tf.math.reduce_mean(pcc))