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

Cudnn issue while using OpenNMT-tf 2.10 with tensorflow 2.2 inAnaconda virtual env

我正在尝试在 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' 页面上尝试 运行 可用的命令时,它在制作 vocab 时识别了 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,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) ]]
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)

这里要是有人指导就好了

Ps。我不认为这是一个版本问题,因为我验证了 openNMT-tf 2.10 需要 tensorflow 2.2 并且安装了 tensorflow-gpu 2.2,anaconda 自己安装了 cuda 10.1 和 cudnn 7.6.5(默认处理 GPU 依赖)。

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

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

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