在 Anaconda 虚拟环境中使用 OpenNMT-tf 2.10 和 tensorflow 2.2 时出现 Cudnn 问题

如何解决在 Anaconda 虚拟环境中使用 OpenNMT-tf 2.10 和 tensorflow 2.2 时出现 Cudnn 问题

我正在尝试在 GPU GeForce RTX 2060 8GB 内存上训练 opennmt-tf 变压器模型。您可以看到步骤 Here。

我已经创建了 Anaconda 虚拟环境并使用以下命令安装了 tensorflow-gpu。

conda install tensorflow-gpu==2.2.0

运行上述命令后，conda env 将处理所有事情，并将在 env 中安装 cuda 10.1 和 cudnn 7.6.5。然后我使用以下命令安装了与 tf 2.2 gpu 兼容的 opennmt-tf 2.10。

~/anaconda3/envs/nmt/bin/pip install opennmt-tf==2.10

以上命令将在 conda 环境中安装 opennmt。

当我尝试运行 opennmt-tf 文档中“Quicstart”页面上可用的命令时，它在制作词汇时识别出了 GPU。但是当我开始训练变压器模型时，它给出了以下 cudnn 错误。

tensorflow.python.framework.errors_impl.InternalError: 2 root error(s) found.
  (0) Internal:  cuDNN launch failure : input shape ([1,504,512,1])
     [[node transformer_base/self_attention_decoder/self_attention_decoder_layer/transformer_layer_wrapper_12/layer_norm_14/FusedBatchnormV3 (defined at /site-packages/opennmt/layers/common.py:128) ]]
     [[Func/gradients/global_norm/write_summary/summary_cond/then/_302/input/_893/_52]]
  (1) Internal:  cuDNN launch failure : input shape ([1,1])
     [[node transformer_base/self_attention_decoder/self_attention_decoder_layer/transformer_layer_wrapper_12/layer_norm_14/FusedBatchnormV3 (defined at /site-packages/opennmt/layers/common.py:128) ]]
0 successful operations.
0 derived errors ignored. [Op:__inference__accumulate_next_33440]

Function call stack:
_accumulate_next -> _accumulate_next

2021-03-01 13:01:01.138811: I tensorflow/stream_executor/stream.cc:1990] [stream=0x560490f17b10,impl=0x560490f172c0] did not wait for [stream=0x5604906de830,impl=0x560490f17250]
2021-03-01 13:01:01.138856: I tensorflow/stream_executor/stream.cc:4938] [stream=0x560490f17b10,impl=0x560490f172c0] did not memcpy host-to-device; source: 0x7ff4467f8780
2021-03-01 13:01:01.138957: F tensorflow/core/common_runtime/gpu/gpu_util.cc:340] cpu->GPU Memcpy Failed
Aborted (core dumped)

如果有人能在这里指导就好了。

附言。我不认为这是一个版本问题，因为我验证了 opennmt-tf 2.10 需要 tensorflow 2.2 并且安装了 tensorflow-gpu 2.2，anaconda 自己安装了 cuda 10.1 和 cudnn 7.6.5（默认处理 GPU 依赖）。

解决方法

这是内存问题。有些人在 StackOverflow 上就 cudnn 问题提出了一些建议。在运行此命令之前，将环境变量 'TF_FORCE_GPU_ALLOW_GROWTH' 设置为 true。

import os
os.environ['TF_FORCE_GPU_ALLOW_GROWTH'] = "true"
os.systme('onmt-main --model_type Transformer --config data.yml train --with_eval')

我终于开始使用上述脚本进行训练，它解决了我的问题。