Sagemaker ml.g4dn.xlarge实例升级TensorFlow至2.4.1后无法识别GPU如何解决

问题根源 #

你手动通过pip升级的TensorFlow 2.4.1与amazonei_tensorflow2_p36内核预装的CUDA、cuDNN版本不匹配。TensorFlow 2.4.1要求运行环境为CUDA 11.0 + cuDNN 8.0，而该内核原生的CUDA版本仅适配预装的TensorFlow 2.3.2，版本对应关系断裂后就会出现GPU识别为空的情况。另外你同时安装tensorflow和tensorflow-gpu两个包也会引发依赖冲突，TensorFlow 2.0之后CPU/GPU支持已经合并到同一个tensorflow包中，不需要单独安装tensorflow-gpu。

修复步骤 #

• 第一步：先卸载现有冲突的TensorFlow相关包
  执行命令：
  !pip uninstall -y tensorflow tensorflow-gpu tensorflow-estimator
• 第二步：从以下两种可行方案中选一种操作即可

  方案一：手动安装适配依赖的TensorFlow 2.4.1 #

  通过conda安装会自动处理CUDA、cuDNN的版本匹配，执行命令：
  !conda install -y tensorflow==2.4.1 cudatoolkit=11.0 cudnn=8.0
  安装完成后必须重启Jupyter内核，再执行GPU检测命令即可正常返回GPU设备信息。

  方案二：使用SageMaker预适配的官方内核 #

  不需要手动修改依赖，直接在Notebook界面的内核切换菜单，选择conda_tensorflow2_p37内核即可。该内核预装TensorFlow 2.4+版本，已完成GPU驱动适配，开箱即可调用GPU资源。

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