为什么这个keras模型会导致内存/显存不足？

我编写了一个keras模型用以进行音频分类，它接受一个音频tensor并输出三个相互独立的标签，代码如下：

wave=keras.Input(shape=(None,))
x=layers.Embedding(256, 64)(wave)
x=layers.LSTM(128)(x)
ft=layers.Dense(64)(x)
ft=layers.Dense(7)(ft)
li_di=layers.concatenate([ft,x])
li=layers.Dense(32)(li_di)
li=layers.Dense(2)(li)
di=layers.Dense(64)(li_di)
di=layers.Dense(5)(di)

model = keras.Model(
    inputs=wave,
    outputs=[ft,li,di],
)

模型结构：

但是当我在colab fit它的时候就会自动因为内存溢出崩溃（CPU模式）或者显存不足（GPU模式）报错...减小数据集或减小embedding的字典大小均不能解决问题。有人遇到过类似的情况或者知道如何解决吗？

附：dataset结构

for i in all_three_train.take(1):
  print(i)

(<tf.Tensor: shape=(385140,), dtype=float32, numpy=
array([ 0.00384521,  0.00387573,  0.00317383, ..., -0.00143433,
       -0.00128174, -0.00134277], dtype=float32)>, <tf.Tensor: shape=(3,), dtype=int32, numpy=array([3, 0, 3], dtype=int32)>)

 

输入管道：

all_three_train=all_three_train.shuffle(1214).padded_batch(32).cache().prefetch(AUTOTUNE)
all_three_test=all_three_test.shuffle(200).padded_batch(32).cache().prefetch(AUTOTUNE)

 

GPU模式下的报错：

CancelledError:  [_Derived_]RecvAsync is cancelled.
     [[{<!-- -->{node gradient_tape/model_2/embedding_2/embedding_lookup/Reshape/_34}}]] [Op:__inference_train_function_5597]

Function call stack:
train_function

 

不知道你这个问题是否已经解决, 如果还没有解决的话:

• 看下这篇博客，也许你就懂了，链接：Keras内置的几种评价函数

如果你已经解决了该问题, 非常希望你能够分享一下解决方案, 写成博客, 将相关链接放在评论区, 以帮助更多的人 ^-^