我在尝试在keras中自定义一个特殊的注意力层。但是尝试了很多方法后，我仍然很困惑为什么总是会出现这个错误。

Traceback (most recent call last): File "D:/Users/LawLi/PyCharmProjects/fixed_talentDNA/adx.py", line 52, in <module>   print(model.predict([tensor1, tensor2, indices]))  # (bs1, sl1, sl2) File "D:\Users\LawLi\Anaconda3\lib\site-packages\keras\engine\training.py", line 1172, in predict   steps=steps) File "D:\Users\LawLi\Anaconda3\lib\site-packages\keras\engine\training_arrays.py", line 293, in predict_loop   ins_batch = slice_arrays(ins, batch_ids) File "D:\Users\LawLi\Anaconda3\lib\site-packages\keras\utils\generic_utils.py", line 507, in slice_arrays   return [None if x is None else x[start] for x in arrays] File "D:\Users\LawLi\Anaconda3\lib\site-packages\keras\utils\generic_utils.py", line 507, in <listcomp>   return [None if x is None else x[start] for x in arrays]IndexError: index 4 is out of bounds for axis 0 with size 4

这是我关于自定义层的测试代码，代码中包含了一个预测的示例。可以直接运行。

from keras.layers import *from keras.models import Modelfrom keras.utils import to_categoricalfrom keras.layers.merge import *class CustomLayer(Layer):    def __init__(self, **kwargs):        self.supports_masking = True        super(CustomLayer, self).__init__(**kwargs)    def build(self, input_shape):        assert len(input_shape) == 3        super(CustomLayer, self).build(input_shape)    def compute_mask(self, inputs, mask=None):        return None    def call(self, x, mask=None):        tensor1, tensor2, ind = x[0], x[1], x[2]  # (bs1, sl1, wd) (bs2, sl2, wd) (bs1, bs2. sl1, sl2)        tensor2 = K.permute_dimensions(tensor2, [0, 2, 1])        align = K.dot(tensor1, tensor2)        align = K.permute_dimensions(align, [0, 2, 1, 3])  # (bs1, bs2, sl1, sl2)        align = align + ind        align = K.max(align, axis=1)        align = K.sum(align, axis=2)        align = K.softmax(align, axis=1)        weighted_ans = tensor1 * K.expand_dims(align, 2)        return K.sum(weighted_ans, axis=1)    def compute_output_shape(self, input_shape):        t1_shape, t2_shape = input_shape[0], input_shape[1]        return t1_shape[0], t1_shape[1], t1_shape[2]# model examplet1 = Input(shape=(7, 3))t2 = Input(batch_shape=(4, 6, 3))t3 = Input(shape=(4, 7, 6))output = CustomLayer()([t1, t2, t3])model = Model([t1, t2, t3], output)# data exampletensor1 = np.random.rand(10, 7, 3)  # (bs1, sl1, wd)tensor2 = np.random.rand(4, 6, 3)  # (bs2, sl2, wd)indices = np.array([0, 1, 3, 2, 0, 1, 2, 2, 3, 1])  # (bs1, 1)indices = to_categorical(indices, num_classes=4) * 999 - 999  # (bs1, bs2)indices = np.expand_dims(indices, axis=2)indices = np.expand_dims(indices, axis=3)indices = np.repeat(indices, 7, axis=2).repeat(6, axis=3)print(model.predict([tensor1, tensor2, indices]))  # (bs1, sl1, wd)

谢谢你的帮助。

回答：

预测函数无法接受第0轴大小不同的输入。因此，我尝试了另一个名为predict_on_batch()的函数