在P100 GPU运行tf_cnn_benchmarks遇TensorFlow报错求助

嘿，这个问题我之前帮朋友排查过类似的，咱们先搞清楚根源，再给你几个可行的解决办法：

为什么指定GPU却涉及CPU操作？ #

你用了--variable_update=parameter_server这个参数，这个变量更新策略原本是为分布式多机多GPU场景设计的。在这种模式下，TensorFlow会把部分变量管理、甚至一些预处理/池化操作默认放到CPU的参数服务器节点上运行，而ResNet50在tf_cnn_benchmarks里默认用的是NCHW数据格式（GPU优化的格式），但CPU的MaxPoolingOp只支持NHWC格式，这就直接触发了报错。

可行的解决方案 #

方案1：更换更适合单GPU的变量更新策略 #

单GPU场景下，parameter_server其实不是最优选择，换成replicated或者distributed_replicated策略，能让几乎所有计算都跑在GPU上，彻底避开CPU格式兼容问题。修改后的命令：

srun python tf_cnn_benchmarks.py --num_gpus=1 --batch_size=32 --model=resnet50 --variable_update=replicated

方案2：强制使用CPU兼容的数据格式（如果一定要用parameter_server） #

如果你因为某些原因必须保留parameter_server策略，可以通过--data_format=NHWC参数强制让模型使用CPU支持的数据格式，这样CPU上的MaxPool操作就能正常运行了：

srun python tf_cnn_benchmarks.py --num_gpus=1 --batch_size=32 --model=resnet50 --variable_update=parameter_server --data_format=NHWC

方案3：考虑升级TensorFlow版本（可选） #

TensorFlow 1.6是比较老的版本（2018年的），而P100是Volta架构的GPU，后续的TF版本（比如1.10+）对Volta架构的兼容性更好，也修复了不少分布式策略下的小bug。如果你的环境允许升级，可以试试这条路径。

内容的提问来源于stack exchange，提问作者Abid Muslim Malik