tensorflow/keras 模型搭建以及第一个epoch速度极慢的问题

问题描述

使用GPU版本的keras搭建深度学习模型时，在模型顺序构造部分【从Sequential()到add()】会卡住一段时间；而在模型开始训练时，会在第一个epoch卡住相对很长的时间之后才能开始训练。

问题分析

在搜索引擎上查找，有不少人提出过相关问题，但是没有找到一个较好的解决方案。折腾了一会儿之后，猜测是 tensorflow / keras 版本的问题。经实践，在conda下创建了一个python3.9的环境，并用conda安装相关包后可以流畅运行（tensorflow2.5版本）。在此基础上继续查找资料，最终认为本质上是cuda和cudnn的问题。conda默认安装的cuda和cudnn版本是跟随环境中的python版本，这就造成了它们与显卡驱动的不匹配，然后出现了上述的情况。在测试环境下，3.9的python版本使得conda在安装tensorflow-gpu的情况下选择了新版本的cuda和cudnn库，使得模型能正常构建与训练。
简而言之，这个问题一般是由cuda和cudnn的版本落后于显卡驱动造成的。例如，30系列的显卡，它们只支持cuda11、cudnn8以上的版本，使用旧版本就会出现这样的问题。

解决方案

请根据自身显卡型号和驱动版本，选择合适的cuda和cudnn，然后再此基础上在选择 tensorflow / keras 的版本。尤其在使用conda之类包管理工具时，不要无脑conda install package，确认完依赖库版本之后，再指定版本，进行包管理。