TF1.x + RTX 3090 训练PWCNet 踩坑
摘要
TensorFlow官方TF1.14~TF1.15.5 不支持RTX3090，TF1.14、TF1.15使用CUDA10训练导致不可避免的NaN问题。使用Nvidia 版本的tensorflow可以支持TF1.x版本和新的硬件如RTX3090。
问题
python版本：python 3.6.13。
Tensorflow 官网安装引导显示只测试了TF1.15.0 使用CUDA10.0 测试过的组合。
直接安装TensorFlow1.14.0、1.15.0、1.15.5后，用下述代码测试，均显示缺少libcudart.so.10.0、libcudnn.so.7等。
import tensorflow as tf
tf.test.is_cpu_available() # 测试CPU是否可用，并显示加载cuda库是否成功
使用conda安装缺少的库：
conda install cudatoolkit=10.0 cudnn
安装后，缺库的警告消失，显示加载cuda库成功，tf.test.is_gpu_available()返回True。
使用该配置训练PWCNet 实现1、实现2均出现NaN和Inf的loss和EPE，甚至在特征计算中输出的特征就是NaN，且每次启动训练后都会在刚开始的几步就出现。数据吞吐量倒正常。
后怀疑TF1.15不支持 RTX 3090。
下载TensorFlow1.x的源码，搜索grep -rn "cudnn.so*" *后只显示cudnn.so.7的关键字，可判定其不支持其他版本的cuda如cuda11、cudnn8等。
随后搜索发现，nvidia提供了TF1.x对RTX 3090、cuda11等新硬件的支持。
I believe this cannot work. Tensorflow 1.15 requires cuda 10:
https://www.tensorflow.org/install/source#gpu
I believe the RTX 3090 does not support cuda 10.
You might need to use either TF 1.15 on CPU. Or a 2.x TF release with a more recent cuda version.
From: https://stackoverflow.com/a/67816052
卸载已有的tensorflow-gpu包和conda安装的cuda包，安装nvidia版本tensorflow:
pip install nvidia-pyindex
pip install nvidia-tensorflow # 会自动安装相关cuda依赖
注：nvidia-tensorflow仓库提示需要使用Python3.8，但我使用Python3.6，可用。
重新执行测试，加载libcudart.so.11.0、libcudnn.so.8成功，GPU可用。
重新启动训练，NaN消失，数据吞吐效率也正常。
使用推荐安装Python3.8后，会自动安装tensorflow1.15.5，最终数据吞吐量略有提升，60samples/sec变为63+samples/sec.
