自定义丢失函数参数无效：begin和size参数应为大小为0的1D张量，但得到的是形状[2]和[2]

Question

我正在尝试将自定义损失函数应用于我的CNN。该损失函数应包含预测（事件）和来自该事件的一些“输入变量”。CNN在输出层输出一个数字，然后“输入变量”在最后一层是concat（这些输入变量因事件而异，用于计算损失）

这一切都很好。然后在损失函数中使用预测和变量，在函数中使用tf.py_函数来计算“实际预测”。这些“真实预测”应用于评估损失。但是，每次运行算法时，都会弹出以下错误：

预期begin和size参数是大小为0的一维张量，但得到了形状[2]和[2]。

我不知道如何修复

以下是损失函数：

            def loss_function(y_true, y_pred):
                def calculate_pt(particle, truth=False):
                    vecFour1 = vectors.LorentzVector()
                    vecFour2 = vectors.LorentzVector()
                    if truth:
                        et1 = np.abs(particle.energy_true[0])/np.cosh(particle.p_eta[0])
                        et2 = np.abs(particle.energy_true[1])/np.cosh(particle.p_eta[1])
                    else:
                        et1 = np.abs(particle.energy[0])/np.cosh(particle.p_eta[0])
                        et2 = np.abs(particle.energy[1])/np.cosh(particle.p_eta[1])
                    vecFour1.setptetaphim(et1, particle.p_eta[0], particle.p_phi[0], 0) #Units in GeV for pt and mass
                    vecFour2.setptetaphim(et2, particle.p_eta[1], particle.p_phi[1], 0)
                    vecFour = vecFour1+vecFour2
                    invM = vecFour.mass
                    return invM
                def func(y_pred):#, y_true, truth=False):
                    y_pred = np.array(y_pred.numpy())
                    df_batch = pd.DataFrame(data = y_pred, columns = ['energy', 'p_e', 'eventNumber', 'p_eta', 'p_phi', 'p_eAccCluster']) # energy_true
                    df_batch = df_batch.groupby("eventNumber").filter(lambda x: len(x) == 2)
                    df_batch = pd.concat([ df_batch, df_batch.shift() ], axis=1)[1::2]
                    mass = df_batch.apply(lambda x: calculate_pt(x), axis=1)
                    mass = np.array(mass.values)
                    return tf.cast(mass, tf.float32)
                print(y_pred)   
                loss_mass = tf.py_function(func=func, inp=[y_pred], Tout=tf.float32, name='loss_function')
                tf.print(loss_mass)
                print('loss1', loss_mass)
                return tf.keras.losses.mse(125, loss_mass)

下面是错误输出：

Compiling!
Compiling done!
Staring training!
Epoch 1/5
Tensor("model_1/concatenate/concat:0", shape=(512, 6), dtype=float32)
loss1 Tensor("loss_function/loss_function:0", dtype=float32, device=/job:localhost/replica:0/task:0)
Tensor("model_1/concatenate/concat:0", shape=(512, 6), dtype=float32)
loss1 Tensor("loss_function/loss_function:0", dtype=float32, device=/job:localhost/replica:0/task:0)
2021-01-23 16:33:25.610533: I tensorflow/compiler/xla/service/service.cc:168] XLA service 0x7f11240b3bd0 initialized for platform CUDA (this does not guarantee that XLA will be used). Devices:
2021-01-23 16:33:25.610566: I tensorflow/compiler/xla/service/service.cc:176]   StreamExecutor device (0): GeForce GTX 1080, Compute Capability 6.1
2021-01-23 16:33:25.686511: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcublas.so.11
2021-01-23 16:33:25.886902: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcublasLt.so.11
2021-01-23 16:33:25.887085: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcudnn.so.8
2021-01-23 16:33:28.392155: I tensorflow/compiler/jit/xla_compilation_cache.cc:333] Compiled cluster using XLA!  This line is logged at most once for the lifetime of the process.
[0 0 0 ... 0 0.122533791 0]
2021-01-23 16:33:28.807058: W tensorflow/core/framework/op_kernel.cc:1763] OP_REQUIRES failed at xla_ops.cc:406 : Invalid argument: Expected begin and size arguments to be 1-D tensors of size 0, but got shapes [2] and [2] instead.
     [[{{node gradient_tape/model_1/concatenate/Slice}}]]
Traceback (most recent call last):
  File "train_model.py", line 310, in <module>
    verbose=1, hparams=None, testing=False)
  File "train_model.py", line 174, in define_and_train_model
    datatype=datatype)
  File "../../deepcalo/model_container.py", line 94, in __init__
    self.train_model()
  File "../../deepcalo/model_container.py", line 610, in train_model
    callbacks=callbacks)
  File "/home/malteal/miniconda3/lib/python3.7/site-packages/tensorflow/python/keras/engine/training.py", line 1100, in fit
    tmp_logs = self.train_function(iterator)
  File "/home/malteal/miniconda3/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py", line 828, in __call__
    result = self._call(*args, **kwds)
  File "/home/malteal/miniconda3/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py", line 888, in _call
    return self._stateless_fn(*args, **kwds)
  File "/home/malteal/miniconda3/lib/python3.7/site-packages/tensorflow/python/eager/function.py", line 2943, in __call__
    filtered_flat_args, captured_inputs=graph_function.captured_inputs)  # pylint: disable=protected-access
  File "/home/malteal/miniconda3/lib/python3.7/site-packages/tensorflow/python/eager/function.py", line 1919, in _call_flat
    ctx, args, cancellation_manager=cancellation_manager))
  File "/home/malteal/miniconda3/lib/python3.7/site-packages/tensorflow/python/eager/function.py", line 560, in call
    ctx=ctx)
  File "/home/malteal/miniconda3/lib/python3.7/site-packages/tensorflow/python/eager/execute.py", line 60, in quick_execute
    inputs, attrs, num_outputs)
tensorflow.python.framework.errors_impl.InvalidArgumentError:  Expected begin and size arguments to be 1-D tensors of size 0, but got shapes [2] and [2] instead.
     [[node gradient_tape/model_1/concatenate/Slice (defined at ../../deepcalo/model_container.py:610) ]]
     [[cluster_1_1/xla_compile]] [Op:__inference_train_function_7443]

Function call stack:
train_function

损失函数只输出一个标量，我认为这很好

希望你们能帮我

提前谢谢