TensorFlow训练机器翻译Transformer时出现FusedBatchNormV3 cuDNN报错

核心报错定位 #

错误触发于TransformerBlock的归一化层调用阶段，核心报错信息为：

tensorflow.python.framework.errors_impl.InternalError: cuDNN launch failure : input shape ([1,4928,256,1]) [Op:FusedBatchNormV3]

具体产生原因 #

• 归一化层类型使用错误
  Transformer架构的标准归一化层为层归一化（Layer Normalization），但你的TransformerBlock中定义的norm1错误使用了批归一化（Batch Normalization），对应报错中的FusedBatchNormV3算子，和DecoderBlock中正确定义的tf.keras.layers.LayerNormalization()逻辑不一致。
• 输入形状不符合批归一化的GPU算子要求
  批归一化的cuDNN实现对输入形状有严格限制，你预期的归一化输入是(batch_size, seq_len, embed_size)格式的三维张量，但实际输入到FusedBatchNormV3的张量为形状[1,4928,256,1]的四维张量：一方面是中间计算过程错误插入了额外维度、或者对序列维度做了异常展平，另一方面批归一化本身不适合NLP序列任务的输入格式，二者叠加触发了cuDNN启动失败。
• 维度参数不匹配
  你提到当前attention和输入相加后的输出形状为(64, 80, 250)，第三个维度为词表大小，但TransformerBlock中配置的嵌入维度为256，和实际输出的250不一致，也会导致张量计算过程中形状异常，最终传入归一化层的张量不符合cuDNN的规格要求。

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