Keras自定义损失函数报错：第二个输入需为标量但形状为[4,1]

问题背景 #

编写Keras自定义损失函数MyLossTennis时触发报错：

tensorflow.python.framework.errors_impl.InvalidArgumentError: Graph execution error: The second input must be a scalar, but it has shape [4,1]

自定义损失函数代码 #

def MyLossTennis(yTrue, yPred):
    noBet = yTrue[:, 0:1]
    favWon = yTrue[:,1:2]
    dogWon = yTrue[:, 2:3]
    
    noBetOdds = yTrue[:, 3:4]
    favOdds = yTrue[:, 4:5]
    dogOdds = yTrue[:, 5:6]
    
    noBetCastka = 0
    favCastka = 1000
    dogCastka = 1000

    def NoBet():
        return K.concatenate([0 * noBet, 0 * noBet, 0 * noBet]) * yPred
    def FavWon():
        return K.concatenate([0 * noBet, favWon * favOdds, -1000 * noBet]) * yPred
    def DogWon():
        return K.concatenate([0 * noBet, -1000 * noBet, dogWon * dogOdds]) * yPred
    
    return (tf.switch_case(tf.cast(tf.argmax([noBet, favWon, dogWon]), tf.int32), branch_fns={0: NoBet, 1: FavWon}, default=DogWon))

完整报错堆栈 #

Epoch 1/100000
2023-03-15 01:09:00.979547: I tensorflow/core/common_runtime/executor.cc:1197] [/job:localhost/replica:0/task:0/device:CPU:0] (DEBUG INFO) Executor start aborting (this does not indicate an error and you can ignore this message): INVALID_ARGUMENT: The second input must be a scalar, but it has shape [4,1]
         [[{{node MyLossTennisPokus/switch_case_1/indexed_case/MyLossTennisPokus/strided_slice_2/_26}}]]
Traceback (most recent call last):
  File "/aux/CreateModels.py", line 875, in <module>
    CreateModel(ModelLolTennis)
  File "/aux/CreateModels.py", line 704, in CreateModel
    FitModel(model, ModelParams, ModelData)
  File "/aux/CreateModels.py", line 157, in FitModel
    model.fit(modelData["xTrain"], modelData["yTrain"], epochs=modelParams["epochs"], verbose=modelParams["testModelVerbose"], batch_size = modelParams["batchSize"], validation_data=(modelData["xVal"], modelData["yVal"]), callbacks=callbacks)
  File "/home/au/.local/lib/python3.10/site-packages/keras/utils/traceback_utils.py", line 70, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/home/au/.local/lib/python3.10/site-packages/tensorflow/python/eager/execute.py", line 52, in quick_execute
    tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
tensorflow.python.framework.errors_impl.InvalidArgumentError: Graph execution error:

The second input must be a scalar, but it has shape [4,1]
         [[{{node MyLossTennisPokus/switch_case_1/indexed_case/MyLossTennisPokus/strided_slice_2/_26}}]] [Op:__inference_train_function_1210]

报错原因 #

1. tf.switch_case参数限制：该函数的第一个参数必须是标量整数张量，用来指定执行哪个分支函数，但你传入的tf.cast(tf.argmax([noBet, favWon, dogWon]), tf.int32)是形状为[4,1]的张量（对应batch中4个样本的索引），不符合标量要求。
2. tf.argmax轴选择错误：noBet、favWon、dogWon都是(batch_size,1)形状的张量，放入列表后tf.argmax默认在轴0计算最大值索引，结果会保留batch维度，无法得到单个标量。

解决方法 #

损失函数需要对batch内每个样本独立计算，不能用tf.switch_case这种单分支控制流，改用向量化操作实现逻辑，避免显式分支判断：

修改后的损失函数代码 #

import tensorflow as tf
from tensorflow import keras as K

def MyLossTennis(yTrue, yPred):
    noBet = yTrue[:, 0:1]
    favWon = yTrue[:,1:2]
    dogWon = yTrue[:, 2:3]
    
    noBetOdds = yTrue[:, 3:4]
    favOdds = yTrue[:, 4:5]
    dogOdds = yTrue[:, 5:6]
    
    # 计算三种场景的损失项
    noBet_loss = K.concatenate([0 * noBet, 0 * noBet, 0 * noBet]) * yPred
    favWon_loss = K.concatenate([0 * noBet, favWon * favOdds, -1000 * noBet]) * yPred
    dogWon_loss = K.concatenate([0 * noBet, -1000 * noBet, dogWon * dogOdds]) * yPred
    
    # 用标签作为掩码，仅保留对应场景的损失
    total_loss = noBet * noBet_loss + favWon * favWon_loss + dogWon * dogWon_loss
    
    # 对每个样本的损失求和，再取batch的平均（符合Keras损失函数输出规范）
    return K.mean(K.sum(total_loss, axis=1))

修改说明 #

• 移除tf.switch_case，直接生成三种场景的损失张量，通过向量化计算覆盖所有样本
• 利用noBet、favWon、dogWon的one-hot属性作为掩码，仅保留当前样本对应场景的损失，其余场景损失被置0
• 最后对batch内的损失取平均，输出标量损失值，符合Keras训练时的损失函数要求

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